diff --git a/ixformer_sdk/.gitignore b/ixformer_sdk/.gitignore new file mode 100644 index 00000000..1691d796 --- /dev/null +++ b/ixformer_sdk/.gitignore @@ -0,0 +1,2 @@ +*.so +build/ diff --git a/ixformer_sdk/__init__.py b/ixformer_sdk/__init__.py new file mode 100644 index 00000000..76f783db --- /dev/null +++ b/ixformer_sdk/__init__.py @@ -0,0 +1,2 @@ +import torch +from .functions import * diff --git a/ixformer_sdk/contrib/DeepCache/__init__.py b/ixformer_sdk/contrib/DeepCache/__init__.py new file mode 100644 index 00000000..2a285982 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/__init__.py @@ -0,0 +1,5 @@ +from .sd.pipeline_stable_diffusion import StableDiffusionPipeline +from .sdxl.pipeline_stable_diffusion_xl import StableDiffusionXLPipeline +from .sdxl.pipeline_stable_diffusion_xl_img2img import StableDiffusionXLImg2ImgPipeline + +from .sd.pipeline_text_to_video_zero import TextToVideoZeroPipeline \ No newline at end of file diff --git a/ixformer_sdk/contrib/DeepCache/ddpm/__init__.py b/ixformer_sdk/contrib/DeepCache/ddpm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/DeepCache/ddpm/ddim.py b/ixformer_sdk/contrib/DeepCache/ddpm/ddim.py new file mode 100644 index 00000000..53213afe --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/ddpm/ddim.py @@ -0,0 +1,180 @@ +import argparse +import traceback +import shutil +import logging +import yaml +import random +import sys +import os +import torch +import numpy as np + +from ddpm.utils.logging import Logger, EmptyLogger +from ddpm.utils.tools import set_random_seed +from accelerate import Accelerator, DistributedDataParallelKwargs + +torch.set_printoptions(sci_mode=False) + +def dict2namespace(config): + namespace = argparse.Namespace() + for key, value in config.items(): + if isinstance(value, dict): + new_value = dict2namespace(value) + else: + new_value = value + setattr(namespace, key, new_value) + return namespace + +def parse_args_and_config(): + parser = argparse.ArgumentParser(description=globals()["__doc__"]) + + parser.add_argument( + "--config", type=str, required=True, help="Path to the config file" + ) + parser.add_argument( + "--seed", type=int, default=1234, help="Random seed") + parser.add_argument( + "--exp", type=str, default="exp", help="Path for saving running related data." + ) + + parser.add_argument( + "--test", action="store_true", help="Whether to test the model" + ) + parser.add_argument( + "--sample", action="store_true", help="Whether to produce samples from the model", + ) + parser.add_argument( + "--image_folder", type=str, default="images", help="folder name for storing the sampled images" + ) + + parser.add_argument( + "--fid", action="store_true" + ) + parser.add_argument( + "--interpolation", action="store_true" + ) + parser.add_argument( + "--resume_training", action="store_true", help="Whether to resume training" + ) + parser.add_argument( + "--ni", action="store_true", help="No interaction. Suitable for Slurm Job launcher", + ) + parser.add_argument( + "--use_pretrained", action="store_true" + ) + parser.add_argument( + "--sample_type", type=str, default="generalized", help="sampling approach (generalized or ddpm_noisy)", + ) + parser.add_argument( + "--skip_type", type=str, default="uniform", help="skip according to (uniform or quadratic)", + ) + parser.add_argument( + "--timesteps", type=int, default=1000, help="number of steps involved" + ) + parser.add_argument( + "--eta", type=float, default=0.0, help="eta used to control the variances of sigma", + ) + parser.add_argument( + "--dyn", action="store_true", help="whether to activate the dynamic train/inference" + ) + parser.add_argument( + "--sequence", action="store_true" + ) + parser.add_argument( + "--select_step", type=int, default=None + ) + parser.add_argument( + "--select_depth", type=int, default=None + ) + + parser.add_argument( + "--cache", action="store_true" + ) + parser.add_argument( + "--cache_interval", type=int, default=None, + ) + parser.add_argument( + "--non_uniform", action="store_true" + ) + parser.add_argument( + "--pow", type=float, default=None, + ) + parser.add_argument( + "--center", type=int, default=None, + ) + parser.add_argument( + "--branch", type=int, default=None, + ) + args = parser.parse_args() + # parse config file + with open(args.config, "r") as f: + config = yaml.safe_load(f) + new_config = dict2namespace(config) + new_config.select_step = args.select_step + new_config.select_depth = args.select_depth + + torch.backends.cudnn.benchmark = True + + return args, new_config + + +def main(): + args, config = parse_args_and_config() + + if args.dyn: + ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + accelerator = Accelerator(kwargs_handlers=[ddp_kwargs]) + else: + accelerator = Accelerator() + args.accelerator = accelerator + + #log_root_dir = "{}_runtime_log".format(args.config[8:-4]) + log_root_dir = "runtime_log" + dataset = args.config[8:-4] + if args.cache: + if args.non_uniform: + sub_dir_name = "{}_{}_cache_{}_pow_{}_center_{}".format(dataset, args.exp, args.cache_interval, args.pow, args.center) + else: + sub_dir_name = "{}_{}_cache_{}".format(dataset, args.exp, args.cache_interval) + else: + sub_dir_name = "{}".format(args.exp) + + if accelerator.is_main_process: + logger = Logger( + root_dir=log_root_dir, + sub_name=sub_dir_name, + config=args.__dict__, + append=(args.sample == True) + ) + args.logger = logger + + args.logger.log("Writing log file to {}".format(args.logger.sub_dir)) + args.logger.log("Exp instance PID = {}".format(os.getpid())) + else: + args.logger = EmptyLogger( + root_dir=log_root_dir, + sub_name=sub_dir_name, + ) + + args.image_folder = args.logger.setup_image_folder("{}".format(args.image_folder)) + + args.seed += accelerator.process_index + # set random seed + set_random_seed(args.seed) + try: + if args.cache: + from ddpm.runners.deepcache import Diffusion + runner = Diffusion(args, config) + runner.sample() + else: + from ddpm.runners.diffusion import Diffusion + runner = Diffusion(args, config) + runner.sample() + except Exception: + logging.error(traceback.format_exc()) + + return 0 + + +if __name__ == "__main__": + main() diff --git a/ixformer_sdk/contrib/DeepCache/ddpm/fid.py b/ixformer_sdk/contrib/DeepCache/ddpm/fid.py new file mode 100644 index 00000000..d04de462 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/ddpm/fid.py @@ -0,0 +1,361 @@ +"""Calculates the Frechet Inception Distance (FID) to evalulate GANs + +The FID metric calculates the distance between two distributions of images. +Typically, we have summary statistics (mean & covariance matrix) of one +of these distributions, while the 2nd distribution is given by a GAN. + +When run as a stand-alone program, it compares the distribution of +images that are stored as PNG/JPEG at a specified location with a +distribution given by summary statistics (in pickle format). + +The FID is calculated by assuming that X_1 and X_2 are the activations of +the pool_3 layer of the inception net for generated samples and real world +samples respectively. + +See --help to see further details. + +Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead +of Tensorflow + +Copyright 2018 Institute of Bioinformatics, JKU Linz + +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 os +import pathlib +from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser + +import numpy as np +import torch +import torchvision.transforms as TF +from PIL import Image +from scipy import linalg +from torch.nn.functional import adaptive_avg_pool2d + +try: + from tqdm import tqdm +except ImportError: + # If tqdm is not available, provide a mock version of it + def tqdm(x): + return x + +from pytorch_fid.inception import InceptionV3 + +parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter) +parser.add_argument('--batch-size', type=int, default=50, + help='Batch size to use') +parser.add_argument('--dataset_name', type=str, default=None) +parser.add_argument('--num-workers', type=int, + help=('Number of processes to use for data loading. ' + 'Defaults to `min(8, num_cpus)`')) +parser.add_argument('--device', type=str, default=None, + help='Device to use. Like cuda, cuda:0 or cpu') +parser.add_argument('--dims', type=int, default=2048, + choices=list(InceptionV3.BLOCK_INDEX_BY_DIM), + help=('Dimensionality of Inception features to use. ' + 'By default, uses pool3 features')) +parser.add_argument('--num_samples', type=int, default=None, + help=('Number of samples for FID estimation')) +parser.add_argument('--res', type=int, default=None, + help=('Resolutions of samples for FID estimation')) +parser.add_argument('--save-stats', action='store_true', + help=('Generate an npz archive from a directory of samples. ' + 'The first path is used as input and the second as output.')) + +parser.add_argument('--path', type=str, nargs=2, + help=('Paths to the generated images or ' + 'to .npz statistic files')) + + +IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm', + 'tif', 'tiff', 'webp'} + + +class ImagePathDataset(torch.utils.data.Dataset): + def __init__(self, files, transforms=None): + self.files = files + self.transforms = transforms + + def __len__(self): + return len(self.files) + + def __getitem__(self, i): + path = self.files[i] + img = Image.open(path).convert('RGB') + if self.transforms is not None: + img = self.transforms(img) + return img + + +def get_activations(files, model, batch_size=50, dims=2048, device='cpu', + num_workers=1, res=None, dataset_name=None): + """Calculates the activations of the pool_3 layer for all images. + + Params: + -- files : List of image files paths + -- model : Instance of inception model + -- batch_size : Batch size of images for the model to process at once. + Make sure that the number of samples is a multiple of + the batch size, otherwise some samples are ignored. This + behavior is retained to match the original FID score + implementation. + -- dims : Dimensionality of features returned by Inception + -- device : Device to run calculations + -- num_workers : Number of parallel dataloader workers + + Returns: + -- A numpy array of dimension (num images, dims) that contains the + activations of the given tensor when feeding inception with the + query tensor. + """ + model.eval() + + if batch_size > len(files): + print(('Warning: batch size is bigger than the data size. ' + 'Setting batch size to data size')) + batch_size = len(files) + + if res is None: + trans = TF.ToTensor() + else: + if dataset_name == 'celeba': + from switchable_diffusion.datasets import Crop + print("In crop image: {}, {}".format(res, dataset_name)) + cx = 89 + cy = 121 + x1 = cy - 64 + x2 = cy + 64 + y1 = cx - 64 + y2 = cx + 64 + trans = TF.Compose([ + Crop(x1, x2, y1, y2), + TF.Resize(res), + TF.ToTensor(), + ]) + else: + trans = TF.Compose([ + TF.Resize(res), + TF.CenterCrop(res), + TF.ToTensor() + ]) + + dataset = ImagePathDataset(files, transforms=trans) + dataloader = torch.utils.data.DataLoader(dataset, + batch_size=batch_size, + shuffle=False, + drop_last=False, + num_workers=num_workers) + + pred_arr = np.empty((len(files), dims)) + + start_idx = 0 + + for batch in tqdm(dataloader): + batch = batch.to(device) + + with torch.no_grad(): + pred = model(batch)[0] + + # If model output is not scalar, apply global spatial average pooling. + # This happens if you choose a dimensionality not equal 2048. + if pred.size(2) != 1 or pred.size(3) != 1: + pred = adaptive_avg_pool2d(pred, output_size=(1, 1)) + + pred = pred.squeeze(3).squeeze(2).cpu().numpy() + + pred_arr[start_idx:start_idx + pred.shape[0]] = pred + + start_idx = start_idx + pred.shape[0] + + return pred_arr + + +def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6): + """Numpy implementation of the Frechet Distance. + The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1) + and X_2 ~ N(mu_2, C_2) is + d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)). + + Stable version by Dougal J. Sutherland. + + Params: + -- mu1 : Numpy array containing the activations of a layer of the + inception net (like returned by the function 'get_predictions') + for generated samples. + -- mu2 : The sample mean over activations, precalculated on an + representative data set. + -- sigma1: The covariance matrix over activations for generated samples. + -- sigma2: The covariance matrix over activations, precalculated on an + representative data set. + + Returns: + -- : The Frechet Distance. + """ + + mu1 = np.atleast_1d(mu1) + mu2 = np.atleast_1d(mu2) + + sigma1 = np.atleast_2d(sigma1) + sigma2 = np.atleast_2d(sigma2) + + assert mu1.shape == mu2.shape, \ + 'Training and test mean vectors have different lengths' + assert sigma1.shape == sigma2.shape, \ + 'Training and test covariances have different dimensions' + + diff = mu1 - mu2 + + # Product might be almost singular + covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) + if not np.isfinite(covmean).all(): + msg = ('fid calculation produces singular product; ' + 'adding %s to diagonal of cov estimates') % eps + print(msg) + offset = np.eye(sigma1.shape[0]) * eps + covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) + + # Numerical error might give slight imaginary component + if np.iscomplexobj(covmean): + if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): + m = np.max(np.abs(covmean.imag)) + raise ValueError('Imaginary component {}'.format(m)) + covmean = covmean.real + + tr_covmean = np.trace(covmean) + + return (diff.dot(diff) + np.trace(sigma1) + + np.trace(sigma2) - 2 * tr_covmean) + + +def calculate_activation_statistics(files, model, batch_size=50, dims=2048, + device='cpu', num_workers=1, res=None, dataset_name=None): + """Calculation of the statistics used by the FID. + Params: + -- files : List of image files paths + -- model : Instance of inception model + -- batch_size : The images numpy array is split into batches with + batch size batch_size. A reasonable batch size + depends on the hardware. + -- dims : Dimensionality of features returned by Inception + -- device : Device to run calculations + -- num_workers : Number of parallel dataloader workers + + Returns: + -- mu : The mean over samples of the activations of the pool_3 layer of + the inception model. + -- sigma : The covariance matrix of the activations of the pool_3 layer of + the inception model. + """ + act = get_activations(files, model, batch_size, dims, device, num_workers, res=res, dataset_name=dataset_name) + mu = np.mean(act, axis=0) + sigma = np.cov(act, rowvar=False) + return mu, sigma + + +def compute_statistics_of_path(path, model, batch_size, dims, device, + num_workers=1, num_samples=None, res=None, dataset_name=None): + if path.endswith('.npz'): + with np.load(path) as f: + m, s = f['mu'][:], f['sigma'][:] + else: + path = pathlib.Path(path) + + files = sorted([file for ext in IMAGE_EXTENSIONS + for file in path.glob('**/*.{}'.format(ext))]) + if num_samples is not None: + #import random + #files = random.sample(files, num_samples) + files = files[:num_samples] + print("Found %d files." % len(files)) + m, s = calculate_activation_statistics(files, model, batch_size, + dims, device, num_workers, res=res, dataset_name=dataset_name) + + return m, s + + +def calculate_fid_given_paths(paths, batch_size, device, dims, num_workers=1, num_samples=None, res=None, dataset_name=None): + """Calculates the FID of two paths""" + for p in paths: + if not os.path.exists(p): + raise RuntimeError('Invalid path: %s' % p) + + block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims] + + model = InceptionV3([block_idx]).to(device) + + m1, s1 = compute_statistics_of_path(paths[0], model, batch_size, + dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name) + m2, s2 = compute_statistics_of_path(paths[1], model, batch_size, + dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name) + fid_value = calculate_frechet_distance(m1, s1, m2, s2) + + return fid_value + + +def save_fid_stats(paths, batch_size, device, dims, num_workers=1, num_samples=None, res=None, dataset_name=None): + """Calculates the FID of two paths""" + if not os.path.exists(paths[0]): + raise RuntimeError('Invalid path: %s' % paths[0]) + + if os.path.exists(paths[1]): + raise RuntimeError('Existing output file: %s' % paths[1]) + + block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims] + + model = InceptionV3([block_idx]).to(device) + + print(f"Saving statistics for {paths[0]}") + + m1, s1 = compute_statistics_of_path(paths[0], model, batch_size, + dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name) + + np.savez_compressed(paths[1], mu=m1, sigma=s1) + + +def main(): + args = parser.parse_args() + + if args.device is None: + device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu') + else: + device = torch.device(args.device) + + if args.num_workers is None: + try: + num_cpus = len(os.sched_getaffinity(0)) + except AttributeError: + # os.sched_getaffinity is not available under Windows, use + # os.cpu_count instead (which may not return the *available* number + # of CPUs). + num_cpus = os.cpu_count() + + num_workers = min(num_cpus, 8) if num_cpus is not None else 0 + else: + num_workers = args.num_workers + + if args.save_stats: + save_fid_stats(args.path, args.batch_size, device, args.dims, num_workers, num_samples=args.num_samples, res=args.res, dataset_name=args.dataset_name) + return + + fid_value = calculate_fid_given_paths(args.path, + args.batch_size, + device, + args.dims, + num_workers, + num_samples=args.num_samples, + res = args.res, dataset_name=args.dataset_name) + print('FID: ', fid_value) + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/ixformer_sdk/contrib/DeepCache/flops.py b/ixformer_sdk/contrib/DeepCache/flops.py new file mode 100644 index 00000000..5efe21be --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/flops.py @@ -0,0 +1,559 @@ +''' +This opcounter is adapted from https://github.com/sovrasov/flops-counter.pytorch and https://github.com/Lyken17/pytorch-OpCounter + +Copyright (C) 2021 Sovrasov V. - All Rights Reserved + * You may use, distribute and modify this code under the + * terms of the MIT license. + * You should have received a copy of the MIT license with + * this file. If not visit https://opensource.org/licenses/MIT +''' +import os +import yaml +import numpy as np +import torch.nn as nn +import torch +has_timm = False + +from diffusers.models.lora import LoRACompatibleLinear, LoRACompatibleConv + +@torch.no_grad() +def count_ops_and_params(model, example_inputs, layer_wise=False): + global CUSTOM_MODULES_MAPPING + ori_model = model + model = copy.deepcopy(model) # deepcopy to avoid changing the original model + flops_model = add_flops_counting_methods(model) + flops_model.eval() + flops_model.start_flops_count(ost=sys.stdout, verbose=False, + ignore_list=[]) + if isinstance(example_inputs, (tuple, list)): + _ = flops_model(*example_inputs) + elif isinstance(example_inputs, dict): + _ = flops_model(**example_inputs) + else: + _ = flops_model(example_inputs) + flops_count, params_count, _layer_flops, _layer_params = flops_model.compute_average_flops_cost() + layer_flops = {} + layer_params = {} + + for m_name, m in model.named_modules(): + layer_flops[m_name] = _layer_flops.get(m) + layer_params[m_name] = _layer_params.get(m) + if layer_wise: + space = 30 - len(m_name) + print("Layer {}: {} MACs = {:.4f} G, Params = {:.4f} M, MACs% = {:.2f}".format( + m_name, ' ' * space, layer_flops[m_name]/1e9, layer_params[m_name] / 1e6, 100 * layer_flops[m_name] / flops_count + )) + + flops_model.stop_flops_count() + CUSTOM_MODULES_MAPPING = {} + #if layer_wise: + # return flops_count, params_count, layer_flops, layer_params + return flops_count, params_count + +def empty_flops_counter_hook(module, input, output): + module.__flops__ += 0 + + +def upsample_flops_counter_hook(module, input, output): + output_size = output[0] + batch_size = output_size.shape[0] + output_elements_count = batch_size + for val in output_size.shape[1:]: + output_elements_count *= val + module.__flops__ += int(output_elements_count) + + +def relu_flops_counter_hook(module, input, output): + active_elements_count = output.numel() + module.__flops__ += int(active_elements_count) + + +def linear_flops_counter_hook(module, input, output): + input = input[0] + # pytorch checks dimensions, so here we don't care much + output_last_dim = output.shape[-1] + bias_flops = output_last_dim if module.bias is not None else 0 + module.__flops__ += int(np.prod(input.shape) * output_last_dim + bias_flops) + + +def pool_flops_counter_hook(module, input, output): + input = input[0] + module.__flops__ += int(np.prod(input.shape)) + + +def bn_flops_counter_hook(module, input, output): + input = input[0] + + batch_flops = np.prod(input.shape) + if module.affine: + batch_flops *= 2 + module.__flops__ += int(batch_flops) + +def ln_flops_counter_hook(module, input, output): + input = input[0] + batch_flops = np.prod(input.shape) + if module.elementwise_affine: + batch_flops *= 2 + module.__flops__ += int(batch_flops) + +def conv_flops_counter_hook(conv_module, input, output): + # Can have multiple inputs, getting the first one + input = input[0] + + batch_size = input.shape[0] + output_dims = list(output.shape[2:]) + + kernel_dims = list(conv_module.kernel_size) + in_channels = conv_module.in_channels + out_channels = conv_module.out_channels + groups = conv_module.groups + + filters_per_channel = out_channels // groups + conv_per_position_flops = int(np.prod(kernel_dims)) * \ + in_channels * filters_per_channel + + active_elements_count = batch_size * int(np.prod(output_dims)) + + overall_conv_flops = conv_per_position_flops * active_elements_count + + bias_flops = 0 + + if conv_module.bias is not None: + + bias_flops = out_channels * active_elements_count + + overall_flops = overall_conv_flops + bias_flops + + conv_module.__flops__ += int(overall_flops) + + +def rnn_flops(flops, rnn_module, w_ih, w_hh, input_size): + # matrix matrix mult ih state and internal state + flops += w_ih.shape[0]*w_ih.shape[1] + # matrix matrix mult hh state and internal state + flops += w_hh.shape[0]*w_hh.shape[1] + if isinstance(rnn_module, (nn.RNN, nn.RNNCell)): + # add both operations + flops += rnn_module.hidden_size + elif isinstance(rnn_module, (nn.GRU, nn.GRUCell)): + # hadamard of r + flops += rnn_module.hidden_size + # adding operations from both states + flops += rnn_module.hidden_size*3 + # last two hadamard product and add + flops += rnn_module.hidden_size*3 + elif isinstance(rnn_module, (nn.LSTM, nn.LSTMCell)): + # adding operations from both states + flops += rnn_module.hidden_size*4 + # two hadamard product and add for C state + flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size + # final hadamard + flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size + return flops + + +def rnn_flops_counter_hook(rnn_module, input, output): + """ + Takes into account batch goes at first position, contrary + to pytorch common rule (but actually it doesn't matter). + If sigmoid and tanh are hard, only a comparison FLOPS should be accurate + """ + flops = 0 + # input is a tuple containing a sequence to process and (optionally) hidden state + inp = input[0] + batch_size = inp[0].shape[0] + seq_length = inp[0].shape[1] + num_layers = rnn_module.num_layers + + for i in range(num_layers): + w_ih = rnn_module.__getattr__('weight_ih_l' + str(i)) + w_hh = rnn_module.__getattr__('weight_hh_l' + str(i)) + if i == 0: + input_size = rnn_module.input_size + else: + input_size = rnn_module.hidden_size + flops = rnn_flops(flops, rnn_module, w_ih, w_hh, input_size) + if rnn_module.bias: + b_ih = rnn_module.__getattr__('bias_ih_l' + str(i)) + b_hh = rnn_module.__getattr__('bias_hh_l' + str(i)) + flops += b_ih.shape[0] + b_hh.shape[0] + + flops *= batch_size + flops *= seq_length + if rnn_module.bidirectional: + flops *= 2 + rnn_module.__flops__ += int(flops) + + +def rnn_cell_flops_counter_hook(rnn_cell_module, input, output): + flops = 0 + inp = input[0] + batch_size = inp.shape[0] + w_ih = rnn_cell_module.__getattr__('weight_ih') + w_hh = rnn_cell_module.__getattr__('weight_hh') + input_size = inp.shape[1] + flops = rnn_flops(flops, rnn_cell_module, w_ih, w_hh, input_size) + if rnn_cell_module.bias: + b_ih = rnn_cell_module.__getattr__('bias_ih') + b_hh = rnn_cell_module.__getattr__('bias_hh') + flops += b_ih.shape[0] + b_hh.shape[0] + + flops *= batch_size + rnn_cell_module.__flops__ += int(flops) + + +def multihead_attention_counter_hook(multihead_attention_module, input, output): + flops = 0 + q, k, v = input + + batch_first = multihead_attention_module.batch_first \ + if hasattr(multihead_attention_module, 'batch_first') else False + if batch_first: + batch_size = q.shape[0] + len_idx = 1 + else: + batch_size = q.shape[1] + len_idx = 0 + + dim_idx = 2 + + qdim = q.shape[dim_idx] + kdim = k.shape[dim_idx] + vdim = v.shape[dim_idx] + + qlen = q.shape[len_idx] + klen = k.shape[len_idx] + vlen = v.shape[len_idx] + + num_heads = multihead_attention_module.num_heads + assert qdim == multihead_attention_module.embed_dim + + if multihead_attention_module.kdim is None: + assert kdim == qdim + if multihead_attention_module.vdim is None: + assert vdim == qdim + + flops = 0 + + # Q scaling + flops += qlen * qdim + # Initial projections + flops += ( + (qlen * qdim * qdim) # QW + + (klen * kdim * kdim) # KW + + (vlen * vdim * vdim) # VW + ) + if multihead_attention_module.in_proj_bias is not None: + flops += (qlen + klen + vlen) * qdim + # attention heads: scale, matmul, softmax, matmul + qk_head_dim = qdim // num_heads + v_head_dim = vdim // num_heads + + head_flops = ( + (qlen * klen * qk_head_dim) # QK^T + + (qlen * klen) # softmax + + (qlen * klen * v_head_dim) # AV + ) + flops += num_heads * head_flops + # final projection, bias is always enabled + flops += qlen * vdim * (vdim + 1) + flops *= batch_size + multihead_attention_module.__flops__ += int(flops) + +def timm_multihead_attention_counter_hook(multihead_attention_module, input, output): + flops = 0 + + q, k, v = input[0], input[0], input[0] + input_dim = input[0].shape[2] + input_len = input[0].shape[1] + batch_size = input[0].shape[0] + + kdim = qdim = vdim = multihead_attention_module.qkv.out_features//3 + qlen = klen = vlen = input_len + + num_heads = multihead_attention_module.num_heads + assert qdim == multihead_attention_module.head_dim * multihead_attention_module.num_heads + + flops = 0 + # Q scaling + flops += qlen * qdim + # Initial projections + flops += ( + (qlen * input_dim * qdim) # QW + + (klen * input_dim * kdim) # KW + + (vlen * input_dim * vdim) # VW + ) + + if multihead_attention_module.qkv.bias is not None: + flops += (qlen + klen + vlen) * qdim + # attention heads: scale, matmul, softmax, matmul + qk_head_dim = qdim // num_heads + v_head_dim = vdim // num_heads + + head_flops = ( + (qlen * klen * qk_head_dim) # QK^T + + (qlen * klen) # softmax + + (qlen * klen * v_head_dim) # AV + ) + flops += num_heads * head_flops + # final projection, bias is always enabled + flops += qlen * vdim * (vdim + 1) + flops *= batch_size + multihead_attention_module.__flops__ += int(flops) + + + +CUSTOM_MODULES_MAPPING = {} + +MODULES_MAPPING = { + # convolutions + nn.Conv1d: conv_flops_counter_hook, + nn.Conv2d: conv_flops_counter_hook, + nn.Conv3d: conv_flops_counter_hook, + LoRACompatibleConv: conv_flops_counter_hook, + # activations + nn.ReLU: relu_flops_counter_hook, + nn.PReLU: relu_flops_counter_hook, + nn.ELU: relu_flops_counter_hook, + nn.LeakyReLU: relu_flops_counter_hook, + nn.ReLU6: relu_flops_counter_hook, + # poolings + nn.MaxPool1d: pool_flops_counter_hook, + nn.AvgPool1d: pool_flops_counter_hook, + nn.AvgPool2d: pool_flops_counter_hook, + nn.MaxPool2d: pool_flops_counter_hook, + nn.MaxPool3d: pool_flops_counter_hook, + nn.AvgPool3d: pool_flops_counter_hook, + nn.AdaptiveMaxPool1d: pool_flops_counter_hook, + nn.AdaptiveAvgPool1d: pool_flops_counter_hook, + nn.AdaptiveMaxPool2d: pool_flops_counter_hook, + nn.AdaptiveAvgPool2d: pool_flops_counter_hook, + nn.AdaptiveMaxPool3d: pool_flops_counter_hook, + nn.AdaptiveAvgPool3d: pool_flops_counter_hook, + # BNs + nn.BatchNorm1d: bn_flops_counter_hook, + nn.BatchNorm2d: bn_flops_counter_hook, + nn.BatchNorm3d: bn_flops_counter_hook, + + nn.InstanceNorm1d: bn_flops_counter_hook, + nn.InstanceNorm2d: bn_flops_counter_hook, + nn.InstanceNorm3d: bn_flops_counter_hook, + nn.GroupNorm: bn_flops_counter_hook, + nn.LayerNorm: ln_flops_counter_hook, + # FC + nn.Linear: linear_flops_counter_hook, + LoRACompatibleLinear: linear_flops_counter_hook, + # Upscale + nn.Upsample: upsample_flops_counter_hook, + # Deconvolution + nn.ConvTranspose1d: conv_flops_counter_hook, + nn.ConvTranspose2d: conv_flops_counter_hook, + nn.ConvTranspose3d: conv_flops_counter_hook, + # RNN + nn.RNN: rnn_flops_counter_hook, + nn.GRU: rnn_flops_counter_hook, + nn.LSTM: rnn_flops_counter_hook, + nn.RNNCell: rnn_cell_flops_counter_hook, + nn.LSTMCell: rnn_cell_flops_counter_hook, + nn.GRUCell: rnn_cell_flops_counter_hook, + nn.MultiheadAttention: multihead_attention_counter_hook +} + +if has_timm: + MODULES_MAPPING.update( + { + timm.models.vision_transformer.Attention: timm_multihead_attention_counter_hook, + } + ) + +if hasattr(nn, 'GELU'): + MODULES_MAPPING[nn.GELU] = relu_flops_counter_hook + + +import sys +from functools import partial +import torch.nn as nn +import copy + +def accumulate_flops(self, layer_flops): + if is_supported_instance(self): + layer_flops[self] = self.__flops__ + return self.__flops__ + else: + sum = 0 + for m in self.children(): + sum += m.accumulate_flops(layer_flops) + layer_flops[self] = sum + return sum + + +def get_model_parameters_number(model): + params_num = sum(p.numel() for p in model.parameters()) + return params_num + + +def add_flops_counting_methods(net_main_module): + # adding additional methods to the existing module object, + # this is done this way so that each function has access to self object + net_main_module.start_flops_count = start_flops_count.__get__(net_main_module) + net_main_module.stop_flops_count = stop_flops_count.__get__(net_main_module) + net_main_module.reset_flops_count = reset_flops_count.__get__(net_main_module) + net_main_module.compute_average_flops_cost = compute_average_flops_cost.__get__( + net_main_module) + + net_main_module.reset_flops_count() + + return net_main_module + +def compute_average_flops_cost(self): + """ + A method that will be available after add_flops_counting_methods() is called + on a desired net object. + Returns current mean flops consumption per image. + """ + + for m in self.modules(): + m.accumulate_flops = accumulate_flops.__get__(m) + + layer_flops = {} + flops_sum = self.accumulate_flops(layer_flops) + + for m in self.modules(): + if hasattr(m, 'accumulate_flops'): + del m.accumulate_flops + + layer_params = {} + for m in self.modules(): + layer_params[m] = get_model_parameters_number(m) + + params_sum = get_model_parameters_number(self) + return flops_sum / self.__batch_counter__, params_sum, layer_flops, layer_params + + +def start_flops_count(self, **kwargs): + """ + A method that will be available after add_flops_counting_methods() is called + on a desired net object. + Activates the computation of mean flops consumption per image. + Call it before you run the network. + """ + add_batch_counter_hook_function(self) + + seen_types = set() + + def add_flops_counter_hook_function(module, ost, verbose, ignore_list): + if type(module) in ignore_list: + seen_types.add(type(module)) + if is_supported_instance(module): + module.__params__ = 0 + elif is_supported_instance(module): + if hasattr(module, '__flops_handle__'): + return + if type(module) in CUSTOM_MODULES_MAPPING: + handle = module.register_forward_hook( + CUSTOM_MODULES_MAPPING[type(module)]) + else: + handle = module.register_forward_hook(MODULES_MAPPING[type(module)]) + module.__flops_handle__ = handle + seen_types.add(type(module)) + else: + if verbose and not type(module) in (nn.Sequential, nn.ModuleList) and \ + not type(module) in seen_types: + print('Warning: module ' + type(module).__name__ + + ' is treated as a zero-op.', file=ost) + seen_types.add(type(module)) + + self.apply(partial(add_flops_counter_hook_function, **kwargs)) + + +def stop_flops_count(self): + """ + A method that will be available after add_flops_counting_methods() is called + on a desired net object. + Stops computing the mean flops consumption per image. + Call whenever you want to pause the computation. + """ + remove_batch_counter_hook_function(self) + self.apply(remove_flops_counter_hook_function) + self.apply(remove_flops_counter_variables) + + +def reset_flops_count(self): + """ + A method that will be available after add_flops_counting_methods() is called + on a desired net object. + Resets statistics computed so far. + """ + add_batch_counter_variables_or_reset(self) + self.apply(add_flops_counter_variable_or_reset) + + +# ---- Internal functions +def batch_counter_hook(module, input, output): + batch_size = 1 + if len(input) > 0: + # Can have multiple inputs, getting the first one + input = input[0] + batch_size = len(input) + else: + pass + print('Warning! No positional inputs found for a module,' + ' assuming batch size is 1.') + module.__batch_counter__ += batch_size + + +def add_batch_counter_variables_or_reset(module): + + module.__batch_counter__ = 0 + + +def add_batch_counter_hook_function(module): + if hasattr(module, '__batch_counter_handle__'): + return + + handle = module.register_forward_hook(batch_counter_hook) + module.__batch_counter_handle__ = handle + + +def remove_batch_counter_hook_function(module): + if hasattr(module, '__batch_counter_handle__'): + module.__batch_counter_handle__.remove() + del module.__batch_counter_handle__ + + +def add_flops_counter_variable_or_reset(module): + if is_supported_instance(module): + if hasattr(module, '__flops__') or hasattr(module, '__params__'): + print('Warning: variables __flops__ or __params__ are already ' + 'defined for the module' + type(module).__name__ + + ' ptflops can affect your code!') + module.__ptflops_backup_flops__ = module.__flops__ + module.__ptflops_backup_params__ = module.__params__ + module.__flops__ = 0 + module.__params__ = get_model_parameters_number(module) + + +def is_supported_instance(module): + if type(module) in MODULES_MAPPING or type(module) in CUSTOM_MODULES_MAPPING: + return True + return False + + +def remove_flops_counter_hook_function(module): + if is_supported_instance(module): + if hasattr(module, '__flops_handle__'): + module.__flops_handle__.remove() + del module.__flops_handle__ + + +def remove_flops_counter_variables(module): + if is_supported_instance(module): + if hasattr(module, '__flops__'): + del module.__flops__ + if hasattr(module, '__ptflops_backup_flops__'): + module.__flops__ = module.__ptflops_backup_flops__ + if hasattr(module, '__params__'): + del module.__params__ + if hasattr(module, '__ptflops_backup_params__'): + module.__params__ = module.__ptflops_backup_params__ + \ No newline at end of file diff --git a/ixformer_sdk/contrib/DeepCache/sd/__init__.py b/ixformer_sdk/contrib/DeepCache/sd/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/DeepCache/sd/pipeline_stable_diffusion.py b/ixformer_sdk/contrib/DeepCache/sd/pipeline_stable_diffusion.py new file mode 100644 index 00000000..578ae1e3 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sd/pipeline_stable_diffusion.py @@ -0,0 +1,812 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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 time +import inspect +from typing import Any, Callable, Dict, List, Optional, Union + +import torch +import numpy as np +from packaging import version +from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer + +from diffusers.configuration_utils import FrozenDict +from diffusers.image_processor import VaeImageProcessor +from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin +from diffusers.models import AutoencoderKL +from diffusers.models.lora import adjust_lora_scale_text_encoder +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils import ( + deprecate, + logging, + replace_example_docstring, +) +from diffusers.utils.torch_utils import randn_tensor + +from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput +from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker + +from .unet_2d_condition import UNet2DConditionModel +from .pipeline_utils import DiffusionPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import StableDiffusionPipeline + + >>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16) + >>> pipe = pipe.to("cuda") + + >>> prompt = "a photo of an astronaut riding a horse on mars" + >>> image = pipe(prompt).images[0] + ``` +""" + +def sample_gaussian_centered(n=1000, sample_size=100, std_dev=100): + samples = [] + + while len(samples) < sample_size: + # Sample from a Gaussian centered at n/2 + sample = int(np.random.normal(loc=n/2, scale=std_dev)) + + # Check if the sample is in bounds + if 1 <= sample < n and sample not in samples: + samples.append(sample) + + return samples + +def sample_from_quad(total_numbers, n_samples, pow=1.2): + while pow > 1: + # Generate linearly spaced values between 0 and a max value + x_values = np.linspace(0, total_numbers**(1/pow), n_samples+1) + + # Raise these values to the power of 1.5 to get a non-linear distribution + indices = np.unique(np.int32(x_values**pow))[:-1] + if len(indices) == n_samples: + break + pow -=0.02 + if pow <= 1: + raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.") + return indices, pow + +def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2): + while pow > 1: + # Generate linearly spaced values between 0 and a max value + x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1) + indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]] + if len(indices) == n_samples: + break + pow -=0.02 + if pow <= 1: + raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.") + return indices, pow + +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + """ + Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + return noise_cfg + + +class StableDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin): + r""" + Pipeline for text-to-image generation using Stable Diffusion. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + The pipeline also inherits the following loading methods: + - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings + - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights + - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights + - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. + text_encoder ([`~transformers.CLIPTextModel`]): + Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). + tokenizer ([`~transformers.CLIPTokenizer`]): + A `CLIPTokenizer` to tokenize text. + unet ([`UNet2DConditionModel`]): + A `UNet2DConditionModel` to denoise the encoded image latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of + [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. + safety_checker ([`StableDiffusionSafetyChecker`]): + Classification module that estimates whether generated images could be considered offensive or harmful. + Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + about a model's potential harms. + feature_extractor ([`~transformers.CLIPImageProcessor`]): + A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. + """ + model_cpu_offload_seq = "text_encoder->unet->vae" + _optional_components = ["safety_checker", "feature_extractor"] + _exclude_from_cpu_offload = ["safety_checker"] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet2DConditionModel, + scheduler: KarrasDiffusionSchedulers, + safety_checker: StableDiffusionSafetyChecker, + feature_extractor: CLIPImageProcessor, + requires_safety_checker: bool = True, + ): + super().__init__() + + if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" + f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " + "to update the config accordingly as leaving `steps_offset` might led to incorrect results" + " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," + " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" + " file" + ) + deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(scheduler.config) + new_config["steps_offset"] = 1 + scheduler._internal_dict = FrozenDict(new_config) + + if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." + " `clip_sample` should be set to False in the configuration file. Please make sure to update the" + " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" + " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" + " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" + ) + deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(scheduler.config) + new_config["clip_sample"] = False + scheduler._internal_dict = FrozenDict(new_config) + + if safety_checker is None and requires_safety_checker: + logger.warning( + f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" + " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" + " results in services or applications open to the public. Both the diffusers team and Hugging Face" + " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" + " it only for use-cases that involve analyzing network behavior or auditing its results. For more" + " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." + ) + + if safety_checker is not None and feature_extractor is None: + raise ValueError( + "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" + " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." + ) + + is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse( + version.parse(unet.config._diffusers_version).base_version + ) < version.parse("0.9.0.dev0") + is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64 + if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64: + deprecation_message = ( + "The configuration file of the unet has set the default `sample_size` to smaller than" + " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the" + " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-" + " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5" + " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the" + " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`" + " in the config might lead to incorrect results in future versions. If you have downloaded this" + " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for" + " the `unet/config.json` file" + ) + deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(unet.config) + new_config["sample_size"] = 64 + unet._internal_dict = FrozenDict(new_config) + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + scheduler=scheduler, + safety_checker=safety_checker, + feature_extractor=feature_extractor, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.register_to_config(requires_safety_checker=requires_safety_checker) + + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.vae.enable_slicing() + + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_slicing() + + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + self.vae.enable_tiling() + + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_tiling() + + def _encode_prompt( + self, + prompt, + device, + num_images_per_prompt, + do_classifier_free_guidance, + negative_prompt=None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + lora_scale: Optional[float] = None, + ): + deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple." + deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False) + + prompt_embeds_tuple = self.encode_prompt( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + do_classifier_free_guidance=do_classifier_free_guidance, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + lora_scale=lora_scale, + ) + + # concatenate for backwards comp + prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]]) + + return prompt_embeds + + def encode_prompt( + self, + prompt, + device, + num_images_per_prompt, + do_classifier_free_guidance, + negative_prompt=None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + lora_scale: Optional[float] = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + do_classifier_free_guidance (`bool`): + whether to use classifier free guidance or not + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + lora_scale (`float`, *optional*): + A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. + """ + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, LoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + adjust_lora_scale_text_encoder(self.text_encoder, lora_scale) + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + # textual inversion: procecss multi-vector tokens if necessary + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, self.tokenizer) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( + text_input_ids, untruncated_ids + ): + removed_text = self.tokenizer.batch_decode( + untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1] + ) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer.model_max_length} tokens: {removed_text}" + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = text_inputs.attention_mask.to(device) + else: + attention_mask = None + + prompt_embeds = self.text_encoder( + text_input_ids.to(device), + attention_mask=attention_mask, + ) + prompt_embeds = prompt_embeds[0] + + if self.text_encoder is not None: + prompt_embeds_dtype = self.text_encoder.dtype + elif self.unet is not None: + prompt_embeds_dtype = self.unet.dtype + else: + prompt_embeds_dtype = prompt_embeds.dtype + + prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) + + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance and negative_prompt_embeds is None: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [""] * batch_size + elif prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = negative_prompt + + # textual inversion: procecss multi-vector tokens if necessary + if isinstance(self, TextualInversionLoaderMixin): + uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer) + + max_length = prompt_embeds.shape[1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = uncond_input.attention_mask.to(device) + else: + attention_mask = None + + negative_prompt_embeds = self.text_encoder( + uncond_input.input_ids.to(device), + attention_mask=attention_mask, + ) + negative_prompt_embeds = negative_prompt_embeds[0] + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) + + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + return prompt_embeds, negative_prompt_embeds + + def run_safety_checker(self, image, device, dtype): + if self.safety_checker is None: + has_nsfw_concept = None + else: + if torch.is_tensor(image): + feature_extractor_input = self.image_processor.postprocess(image, output_type="pil") + else: + feature_extractor_input = self.image_processor.numpy_to_pil(image) + safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device) + image, has_nsfw_concept = self.safety_checker( + images=image, clip_input=safety_checker_input.pixel_values.to(dtype) + ) + return image, has_nsfw_concept + + def decode_latents(self, latents): + deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead" + deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False) + + latents = 1 / self.vae.config.scaling_factor * latents + image = self.vae.decode(latents, return_dict=False)[0] + image = (image / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 + image = image.cpu().permute(0, 2, 3, 1).float().numpy() + return image + + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs( + self, + prompt, + height, + width, + callback_steps, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + ): + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): + shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + return latents + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: int = 1, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guidance_rescale: float = 0.0, + cache_interval: int = 1, + cache_layer_id: int = None, + cache_block_id: int = None, + uniform: bool = True, + pow: float = None, + center: int = None, + output_all_sequence: bool = False, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. + height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): + The width in pixels of the generated image. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, *optional*, defaults to 7.5): + A higher guidance scale value encourages the model to generate images closely linked to the text + `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide what to not include in image generation. If not defined, you need to + pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies + to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If + not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a + plain tuple. + callback (`Callable`, *optional*): + A function that calls every `callback_steps` steps during inference. The function is called with the + following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function is called. If not specified, the callback is called at + every step. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in + [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + guidance_rescale (`float`, *optional*, defaults to 0.7): + Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when + using zero terminal SNR. + + Examples: + + Returns: + [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, + otherwise a `tuple` is returned where the first element is a list with the generated images and the + second element is a list of `bool`s indicating whether the corresponding generated image contains + "not-safe-for-work" (nsfw) content. + """ + # 0. Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds + ) + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + text_encoder_lora_scale = ( + cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None + ) + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt, + device, + num_images_per_prompt, + do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + lora_scale=text_encoder_lora_scale, + ) + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + if do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + + # 4. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # 5. Prepare latent variables + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + + prv_features = None + latents_list = [latents] + + if cache_interval == 1: + interval_seq = list(range(num_inference_steps)) + else: + if uniform: + interval_seq = list(range(0, num_inference_steps, cache_interval)) + else: + num_slow_step = num_inference_steps//cache_interval + if num_inference_steps%cache_interval != 0: + num_slow_step += 1 + + interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,] + #interval_seq, pow = sample_from_quad(num_inference_steps, num_inference_steps//cache_interval, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,] + + interval_seq = sorted(interval_seq) + #print(interval_seq, len(interval_seq), pow) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + #print("[INFO] Update Feature Interval = {}, Update Layer Number = {}, Update Block Number = {}".format(cache_interval, cache_layer_id, cache_block_id)) + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + if i in interval_seq: + prv_features = None + + # predict the noise residual + noise_pred, prv_features = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + replicate_prv_feature=prv_features, + quick_replicate= cache_interval>1, + cache_layer_id=cache_layer_id, + cache_block_id=cache_block_id, + return_dict=False, + ) + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + if do_classifier_free_guidance and guidance_rescale > 0.0: + # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf + noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + latents_list.append(latents) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + callback(i, t, latents) + + if not output_type == "latent": + if output_all_sequence: + image = [self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] for latents in latents_list] + has_nsfw_concept = None #self.run_safety_checker(images[0], device, prompt_embeds.dtype) + num_img = len(image) + else: + image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] + has_nsfw_concept = None + num_img = image.shape[0] + else: + image = latents + has_nsfw_concept = None + + if has_nsfw_concept is None: + do_denormalize = [True] * num_img + else: + do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] + + if output_all_sequence: + image = [self.image_processor.postprocess(img, output_type=output_type, do_denormalize=do_denormalize) for img in image] + else: + image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + + # Offload all models + self.maybe_free_model_hooks() + if not return_dict: + return (image, has_nsfw_concept,) + + return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) \ No newline at end of file diff --git a/ixformer_sdk/contrib/DeepCache/sd/pipeline_text_to_video_zero.py b/ixformer_sdk/contrib/DeepCache/sd/pipeline_text_to_video_zero.py new file mode 100644 index 00000000..1d550307 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sd/pipeline_text_to_video_zero.py @@ -0,0 +1,741 @@ +import copy +from dataclasses import dataclass +from typing import Callable, List, Optional, Union + +import numpy as np +import PIL.Image +import torch +import torch.nn.functional as F +from torch.nn.functional import grid_sample +from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer + +from diffusers.models import AutoencoderKL +from .unet_2d_condition import UNet2DConditionModel +from .pipeline_stable_diffusion import StableDiffusionPipeline, StableDiffusionSafetyChecker +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils import BaseOutput +from diffusers.utils.torch_utils import randn_tensor + +def sample_gaussian_centered(n=1000, sample_size=100, std_dev=100): + samples = [] + + while len(samples) < sample_size: + # Sample from a Gaussian centered at n/2 + sample = int(np.random.normal(loc=n/2, scale=std_dev)) + + # Check if the sample is in bounds + if 1 <= sample < n and sample not in samples: + samples.append(sample) + + return samples + +def sample_from_quad(total_numbers, n_samples, pow=1.2): + while pow > 1: + # Generate linearly spaced values between 0 and a max value + x_values = np.linspace(0, total_numbers**(1/pow), n_samples+1) + + # Raise these values to the power of 1.5 to get a non-linear distribution + indices = np.unique(np.int32(x_values**pow))[:-1] + if len(indices) == n_samples: + break + pow -=0.02 + if pow <= 1: + raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.") + return indices, pow + +def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2): + while pow > 1: + # Generate linearly spaced values between 0 and a max value + x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1) + indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]] + if len(indices) == n_samples: + break + pow -=0.02 + if pow <= 1: + raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.") + return indices, pow + +def rearrange_0(tensor, f): + F, C, H, W = tensor.size() + tensor = torch.permute(torch.reshape(tensor, (F // f, f, C, H, W)), (0, 2, 1, 3, 4)) + return tensor + + +def rearrange_1(tensor): + B, C, F, H, W = tensor.size() + return torch.reshape(torch.permute(tensor, (0, 2, 1, 3, 4)), (B * F, C, H, W)) + + +def rearrange_3(tensor, f): + F, D, C = tensor.size() + return torch.reshape(tensor, (F // f, f, D, C)) + + +def rearrange_4(tensor): + B, F, D, C = tensor.size() + return torch.reshape(tensor, (B * F, D, C)) + + +class CrossFrameAttnProcessor: + """ + Cross frame attention processor. Each frame attends the first frame. + + Args: + batch_size: The number that represents actual batch size, other than the frames. + For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to + 2, due to classifier-free guidance. + """ + + def __init__(self, batch_size=2): + self.batch_size = batch_size + + def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None): + batch_size, sequence_length, _ = hidden_states.shape + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + query = attn.to_q(hidden_states) + + is_cross_attention = encoder_hidden_states is not None + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + elif attn.norm_cross: + encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + # Cross Frame Attention + if not is_cross_attention: + video_length = key.size()[0] // self.batch_size + first_frame_index = [0] * video_length + + # rearrange keys to have batch and frames in the 1st and 2nd dims respectively + key = rearrange_3(key, video_length) + key = key[:, first_frame_index] + # rearrange values to have batch and frames in the 1st and 2nd dims respectively + value = rearrange_3(value, video_length) + value = value[:, first_frame_index] + + # rearrange back to original shape + key = rearrange_4(key) + value = rearrange_4(value) + + query = attn.head_to_batch_dim(query) + key = attn.head_to_batch_dim(key) + value = attn.head_to_batch_dim(value) + + attention_probs = attn.get_attention_scores(query, key, attention_mask) + hidden_states = torch.bmm(attention_probs, value) + hidden_states = attn.batch_to_head_dim(hidden_states) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + return hidden_states + + +class CrossFrameAttnProcessor2_0: + """ + Cross frame attention processor with scaled_dot_product attention of Pytorch 2.0. + + Args: + batch_size: The number that represents actual batch size, other than the frames. + For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to + 2, due to classifier-free guidance. + """ + + def __init__(self, batch_size=2): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + self.batch_size = batch_size + + def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None): + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + inner_dim = hidden_states.shape[-1] + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + query = attn.to_q(hidden_states) + + is_cross_attention = encoder_hidden_states is not None + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + elif attn.norm_cross: + encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + # Cross Frame Attention + if not is_cross_attention: + video_length = max(1, key.size()[0] // self.batch_size) + first_frame_index = [0] * video_length + + # rearrange keys to have batch and frames in the 1st and 2nd dims respectively + key = rearrange_3(key, video_length) + key = key[:, first_frame_index] + # rearrange values to have batch and frames in the 1st and 2nd dims respectively + value = rearrange_3(value, video_length) + value = value[:, first_frame_index] + + # rearrange back to original shape + key = rearrange_4(key) + value = rearrange_4(value) + + head_dim = inner_dim // attn.heads + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # the output of sdp = (batch, num_heads, seq_len, head_dim) + # TODO: add support for attn.scale when we move to Torch 2.1 + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +@dataclass +class TextToVideoPipelineOutput(BaseOutput): + r""" + Output class for zero-shot text-to-video pipeline. + + Args: + images (`[List[PIL.Image.Image]`, `np.ndarray`]): + List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + num_channels)`. + nsfw_content_detected (`[List[bool]]`): + List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or + `None` if safety checking could not be performed. + """ + + images: Union[List[PIL.Image.Image], np.ndarray] + nsfw_content_detected: Optional[List[bool]] + + +def coords_grid(batch, ht, wd, device): + # Adapted from https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py + coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device)) + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch, 1, 1, 1) + + +def warp_single_latent(latent, reference_flow): + """ + Warp latent of a single frame with given flow + + Args: + latent: latent code of a single frame + reference_flow: flow which to warp the latent with + + Returns: + warped: warped latent + """ + _, _, H, W = reference_flow.size() + _, _, h, w = latent.size() + coords0 = coords_grid(1, H, W, device=latent.device).to(latent.dtype) + + coords_t0 = coords0 + reference_flow + coords_t0[:, 0] /= W + coords_t0[:, 1] /= H + + coords_t0 = coords_t0 * 2.0 - 1.0 + coords_t0 = F.interpolate(coords_t0, size=(h, w), mode="bilinear") + coords_t0 = torch.permute(coords_t0, (0, 2, 3, 1)) + + warped = grid_sample(latent, coords_t0, mode="nearest", padding_mode="reflection") + return warped + + +def create_motion_field(motion_field_strength_x, motion_field_strength_y, frame_ids, device, dtype): + """ + Create translation motion field + + Args: + motion_field_strength_x: motion strength along x-axis + motion_field_strength_y: motion strength along y-axis + frame_ids: indexes of the frames the latents of which are being processed. + This is needed when we perform chunk-by-chunk inference + device: device + dtype: dtype + + Returns: + + """ + seq_length = len(frame_ids) + reference_flow = torch.zeros((seq_length, 2, 512, 512), device=device, dtype=dtype) + for fr_idx in range(seq_length): + reference_flow[fr_idx, 0, :, :] = motion_field_strength_x * (frame_ids[fr_idx]) + reference_flow[fr_idx, 1, :, :] = motion_field_strength_y * (frame_ids[fr_idx]) + return reference_flow + + +def create_motion_field_and_warp_latents(motion_field_strength_x, motion_field_strength_y, frame_ids, latents): + """ + Creates translation motion and warps the latents accordingly + + Args: + motion_field_strength_x: motion strength along x-axis + motion_field_strength_y: motion strength along y-axis + frame_ids: indexes of the frames the latents of which are being processed. + This is needed when we perform chunk-by-chunk inference + latents: latent codes of frames + + Returns: + warped_latents: warped latents + """ + motion_field = create_motion_field( + motion_field_strength_x=motion_field_strength_x, + motion_field_strength_y=motion_field_strength_y, + frame_ids=frame_ids, + device=latents.device, + dtype=latents.dtype, + ) + warped_latents = latents.clone().detach() + for i in range(len(warped_latents)): + warped_latents[i] = warp_single_latent(latents[i][None], motion_field[i][None]) + return warped_latents + + +class TextToVideoZeroPipeline(StableDiffusionPipeline): + r""" + Pipeline for zero-shot text-to-video generation using Stable Diffusion. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`CLIPTextModel`]): + Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). + tokenizer (`CLIPTokenizer`): + A [`~transformers.CLIPTokenizer`] to tokenize text. + unet ([`UNet2DConditionModel`]): + A [`UNet3DConditionModel`] to denoise the encoded video latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of + [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. + safety_checker ([`StableDiffusionSafetyChecker`]): + Classification module that estimates whether generated images could be considered offensive or harmful. + Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details + about a model's potential harms. + feature_extractor ([`CLIPImageProcessor`]): + A [`CLIPImageProcessor`] to extract features from generated images; used as inputs to the `safety_checker`. + """ + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet2DConditionModel, + scheduler: KarrasDiffusionSchedulers, + safety_checker: StableDiffusionSafetyChecker, + feature_extractor: CLIPImageProcessor, + requires_safety_checker: bool = True, + ): + super().__init__( + vae, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker + ) + processor = ( + CrossFrameAttnProcessor2_0(batch_size=2) + if hasattr(F, "scaled_dot_product_attention") + else CrossFrameAttnProcessor(batch_size=2) + ) + self.unet.set_attn_processor(processor) + + def forward_loop(self, x_t0, t0, t1, generator): + """ + Perform DDPM forward process from time t0 to t1. This is the same as adding noise with corresponding variance. + + Args: + x_t0: + Latent code at time t0. + t0: + Timestep at t0. + t1: + Timestamp at t1. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + + Returns: + x_t1: + Forward process applied to x_t0 from time t0 to t1. + """ + eps = randn_tensor(x_t0.size(), generator=generator, dtype=x_t0.dtype, device=x_t0.device) + alpha_vec = torch.prod(self.scheduler.alphas[t0:t1]) + x_t1 = torch.sqrt(alpha_vec) * x_t0 + torch.sqrt(1 - alpha_vec) * eps + return x_t1 + + def backward_loop( + self, + latents, + timesteps, + prompt_embeds, + guidance_scale, + callback, + callback_steps, + num_warmup_steps, + extra_step_kwargs, + prv_features, + interval_seq, + cache_interval, + cache_block_id, + cache_layer_id, + cross_attention_kwargs=None, + ): + """ + Perform backward process given list of time steps. + + Args: + latents: + Latents at time timesteps[0]. + timesteps: + Time steps along which to perform backward process. + prompt_embeds: + Pre-generated text embeddings. + guidance_scale: + A higher guidance scale value encourages the model to generate images closely linked to the text + `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. + callback (`Callable`, *optional*): + A function that calls every `callback_steps` steps during inference. The function is called with the + following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function is called. If not specified, the callback is called at + every step. + extra_step_kwargs: + Extra_step_kwargs. + cross_attention_kwargs: + A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in + [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + num_warmup_steps: + number of warmup steps. + + Returns: + latents: + Latents of backward process output at time timesteps[-1]. + """ + do_classifier_free_guidance = guidance_scale > 1.0 + num_steps = (len(timesteps) - num_warmup_steps) // self.scheduler.order + with self.progress_bar(total=num_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + ######## + if i in interval_seq: + prv_features = None + # predict the noise residual + noise_pred, prv_features = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + replicate_prv_feature=prv_features, + quick_replicate= cache_interval>1, + cache_layer_id=cache_layer_id, + cache_block_id=cache_block_id, + return_dict=False, + ) + ######## + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + return latents.clone().detach() + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + video_length: Optional[int] = 8, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_videos_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + motion_field_strength_x: float = 12, + motion_field_strength_y: float = 12, + output_type: Optional[str] = "tensor", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + t0: int = 44, + t1: int = 47, + frame_ids: Optional[List[int]] = None, + ######## + cache_interval: int = 1, + cache_layer_id: int = None, + cache_block_id: int = None, + uniform: bool = True, + pow: float = None, + center: int = None, + output_all_sequence: bool = False, + ######## + ): + """ + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. + video_length (`int`, *optional*, defaults to 8): + The number of generated video frames. + height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): + The width in pixels of the generated image. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, *optional*, defaults to 7.5): + A higher guidance scale value encourages the model to generate images closely linked to the text + `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide what to not include in video generation. If not defined, you need to + pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies + to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for video + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + output_type (`str`, *optional*, defaults to `"numpy"`): + The output format of the generated video. Choose between `"latent"` and `"numpy"`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a + [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`] instead of + a plain tuple. + callback (`Callable`, *optional*): + A function that calls every `callback_steps` steps during inference. The function is called with the + following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function is called. If not specified, the callback is called at + every step. + motion_field_strength_x (`float`, *optional*, defaults to 12): + Strength of motion in generated video along x-axis. See the [paper](https://arxiv.org/abs/2303.13439), + Sect. 3.3.1. + motion_field_strength_y (`float`, *optional*, defaults to 12): + Strength of motion in generated video along y-axis. See the [paper](https://arxiv.org/abs/2303.13439), + Sect. 3.3.1. + t0 (`int`, *optional*, defaults to 44): + Timestep t0. Should be in the range [0, num_inference_steps - 1]. See the + [paper](https://arxiv.org/abs/2303.13439), Sect. 3.3.1. + t1 (`int`, *optional*, defaults to 47): + Timestep t0. Should be in the range [t0 + 1, num_inference_steps - 1]. See the + [paper](https://arxiv.org/abs/2303.13439), Sect. 3.3.1. + frame_ids (`List[int]`, *optional*): + Indexes of the frames that are being generated. This is used when generating longer videos + chunk-by-chunk. + + Returns: + [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`]: + The output contains a `ndarray` of the generated video, when `output_type` != `"latent"`, otherwise a + latent code of generated videos and a list of `bool`s indicating whether the corresponding generated + video contains "not-safe-for-work" (nsfw) content.. + """ + assert video_length > 0 + if frame_ids is None: + frame_ids = list(range(video_length)) + assert len(frame_ids) == video_length + + assert num_videos_per_prompt == 1 + + if isinstance(prompt, str): + prompt = [prompt] + if isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] + + # Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # Check inputs. Raise error if not correct + self.check_inputs(prompt, height, width, callback_steps) + + # Define call parameters + batch_size = 1 if isinstance(prompt, str) else len(prompt) + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # Encode input prompt + prompt_embeds = self._encode_prompt( + prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt + ) + + # Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # Prepare latent variables + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + # Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + + prv_features = None #record cache feature **** + latents_list = [latents] + + if cache_interval == 1: + interval_seq = list(range(num_inference_steps)) + else: + if uniform: + interval_seq = list(range(0, num_inference_steps, cache_interval)) + else: + num_slow_step = num_inference_steps//cache_interval + if num_inference_steps%cache_interval != 0: + num_slow_step += 1 + + interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,] + #interval_seq, pow = sample_from_quad(num_inference_steps, num_inference_steps//cache_interval, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,] + + interval_seq = sorted(interval_seq) + + # Perform the first backward process up to time T_1 + x_1_t1 = self.backward_loop( + timesteps=timesteps[: -t1 - 1], + prompt_embeds=prompt_embeds, + latents=latents, + guidance_scale=guidance_scale, + callback=callback, + callback_steps=callback_steps, + extra_step_kwargs=extra_step_kwargs, + num_warmup_steps=num_warmup_steps, + prv_features=prv_features, + interval_seq=interval_seq, + cache_interval=cache_interval, + cache_block_id=cache_block_id, + cache_layer_id=cache_layer_id, + ) + scheduler_copy = copy.deepcopy(self.scheduler) + + # Perform the second backward process up to time T_0 + x_1_t0 = self.backward_loop( + timesteps=timesteps[-t1 - 1 : -t0 - 1], + prompt_embeds=prompt_embeds, + latents=x_1_t1, + guidance_scale=guidance_scale, + callback=callback, + callback_steps=callback_steps, + extra_step_kwargs=extra_step_kwargs, + num_warmup_steps=0, + prv_features=prv_features, + interval_seq=interval_seq, + cache_interval=cache_interval, + cache_block_id=cache_block_id, + cache_layer_id=cache_layer_id, + ) + + # Propagate first frame latents at time T_0 to remaining frames + x_2k_t0 = x_1_t0.repeat(video_length - 1, 1, 1, 1) + + # Add motion in latents at time T_0 + x_2k_t0 = create_motion_field_and_warp_latents( + motion_field_strength_x=motion_field_strength_x, + motion_field_strength_y=motion_field_strength_y, + latents=x_2k_t0, + frame_ids=frame_ids[1:], + ) + + # Perform forward process up to time T_1 + x_2k_t1 = self.forward_loop( + x_t0=x_2k_t0, + t0=timesteps[-t0 - 1].item(), + t1=timesteps[-t1 - 1].item(), + generator=generator, + ) + + # Perform backward process from time T_1 to 0 + x_1k_t1 = torch.cat([x_1_t1, x_2k_t1]) + b, l, d = prompt_embeds.size() + prompt_embeds = prompt_embeds[:, None].repeat(1, video_length, 1, 1).reshape(b * video_length, l, d) + + self.scheduler = scheduler_copy + x_1k_0 = self.backward_loop( + timesteps=timesteps[-t1 - 1 :], + prompt_embeds=prompt_embeds, + latents=x_1k_t1, + guidance_scale=guidance_scale, + callback=callback, + callback_steps=callback_steps, + extra_step_kwargs=extra_step_kwargs, + num_warmup_steps=0, + prv_features=prv_features, + interval_seq=interval_seq, + cache_interval=cache_interval, + cache_block_id=cache_block_id, + cache_layer_id=cache_layer_id, + ) + latents = x_1k_0 + + # manually for max memory savings + if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: + self.unet.to("cpu") + torch.cuda.empty_cache() + + if output_type == "latent": + image = latents + has_nsfw_concept = None + else: + image = self.decode_latents(latents) + # Run safety checker + image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image, has_nsfw_concept) + + return TextToVideoPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) \ No newline at end of file diff --git a/ixformer_sdk/contrib/DeepCache/sd/pipeline_utils.py b/ixformer_sdk/contrib/DeepCache/sd/pipeline_utils.py new file mode 100644 index 00000000..fc0bef92 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sd/pipeline_utils.py @@ -0,0 +1,1839 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# +# 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 fnmatch +import importlib +import inspect +import os +import re +import sys +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import PIL +import torch +from huggingface_hub import ModelCard, create_repo, hf_hub_download, model_info, snapshot_download +from packaging import version +from requests.exceptions import HTTPError +from tqdm.auto import tqdm + +import diffusers + +from diffusers import __version__ +from diffusers.configuration_utils import ConfigMixin +from diffusers.models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT +from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME +from diffusers.utils import ( + CONFIG_NAME, + DEPRECATED_REVISION_ARGS, + # DIFFUSERS_CACHE, + # HF_HUB_OFFLINE, + SAFETENSORS_WEIGHTS_NAME, + WEIGHTS_NAME, + BaseOutput, + deprecate, + get_class_from_dynamic_module, + is_accelerate_available, + is_accelerate_version, + is_torch_version, + is_transformers_available, + logging, + numpy_to_pil, +) +from diffusers.utils.torch_utils import is_compiled_module +from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE +DIFFUSERS_CACHE=HUGGINGFACE_HUB_CACHE +ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} +HF_HUB_OFFLINE = os.getenv("HF_HUB_OFFLINE", "").upper() in ENV_VARS_TRUE_VALUES +if is_transformers_available(): + import transformers + from transformers import PreTrainedModel + from transformers.utils import FLAX_WEIGHTS_NAME as TRANSFORMERS_FLAX_WEIGHTS_NAME + from transformers.utils import SAFE_WEIGHTS_NAME as TRANSFORMERS_SAFE_WEIGHTS_NAME + from transformers.utils import WEIGHTS_NAME as TRANSFORMERS_WEIGHTS_NAME + +from diffusers.utils import FLAX_WEIGHTS_NAME, ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, PushToHubMixin + + +if is_accelerate_available(): + import accelerate + + +INDEX_FILE = "diffusion_pytorch_model.bin" +CUSTOM_PIPELINE_FILE_NAME = "pipeline.py" +DUMMY_MODULES_FOLDER = "diffusers.utils" +TRANSFORMERS_DUMMY_MODULES_FOLDER = "transformers.utils" +CONNECTED_PIPES_KEYS = ["prior"] + + +logger = logging.get_logger(__name__) + + +LOADABLE_CLASSES = { + "diffusers": { + "ModelMixin": ["save_pretrained", "from_pretrained"], + "SchedulerMixin": ["save_pretrained", "from_pretrained"], + "DiffusionPipeline": ["save_pretrained", "from_pretrained"], + "OnnxRuntimeModel": ["save_pretrained", "from_pretrained"], + }, + "transformers": { + "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"], + "PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"], + "PreTrainedModel": ["save_pretrained", "from_pretrained"], + "FeatureExtractionMixin": ["save_pretrained", "from_pretrained"], + "ProcessorMixin": ["save_pretrained", "from_pretrained"], + "ImageProcessingMixin": ["save_pretrained", "from_pretrained"], + }, + "onnxruntime.training": { + "ORTModule": ["save_pretrained", "from_pretrained"], + }, +} + +ALL_IMPORTABLE_CLASSES = {} +for library in LOADABLE_CLASSES: + ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library]) + + +@dataclass +class ImagePipelineOutput(BaseOutput): + """ + Output class for image pipelines. + + Args: + images (`List[PIL.Image.Image]` or `np.ndarray`) + List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + num_channels)`. + """ + + images: Union[List[PIL.Image.Image], np.ndarray] + + +@dataclass +class AudioPipelineOutput(BaseOutput): + """ + Output class for audio pipelines. + + Args: + audios (`np.ndarray`) + List of denoised audio samples of a NumPy array of shape `(batch_size, num_channels, sample_rate)`. + """ + + audios: np.ndarray + + +def is_safetensors_compatible(filenames, variant=None, passed_components=None) -> bool: + """ + Checking for safetensors compatibility: + - By default, all models are saved with the default pytorch serialization, so we use the list of default pytorch + files to know which safetensors files are needed. + - The model is safetensors compatible only if there is a matching safetensors file for every default pytorch file. + + Converting default pytorch serialized filenames to safetensors serialized filenames: + - For models from the diffusers library, just replace the ".bin" extension with ".safetensors" + - For models from the transformers library, the filename changes from "pytorch_model" to "model", and the ".bin" + extension is replaced with ".safetensors" + """ + pt_filenames = [] + + sf_filenames = set() + + passed_components = passed_components or [] + + for filename in filenames: + _, extension = os.path.splitext(filename) + + if len(filename.split("/")) == 2 and filename.split("/")[0] in passed_components: + continue + + if extension == ".bin": + pt_filenames.append(filename) + elif extension == ".safetensors": + sf_filenames.add(filename) + + for filename in pt_filenames: + # filename = 'foo/bar/baz.bam' -> path = 'foo/bar', filename = 'baz', extention = '.bam' + path, filename = os.path.split(filename) + filename, extension = os.path.splitext(filename) + + if filename.startswith("pytorch_model"): + filename = filename.replace("pytorch_model", "model") + else: + filename = filename + + expected_sf_filename = os.path.join(path, filename) + expected_sf_filename = f"{expected_sf_filename}.safetensors" + + if expected_sf_filename not in sf_filenames: + logger.warning(f"{expected_sf_filename} not found") + return False + + return True + + +def variant_compatible_siblings(filenames, variant=None) -> Union[List[os.PathLike], str]: + weight_names = [ + WEIGHTS_NAME, + SAFETENSORS_WEIGHTS_NAME, + FLAX_WEIGHTS_NAME, + ONNX_WEIGHTS_NAME, + ONNX_EXTERNAL_WEIGHTS_NAME, + ] + + if is_transformers_available(): + weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME, TRANSFORMERS_FLAX_WEIGHTS_NAME] + + # model_pytorch, diffusion_model_pytorch, ... + weight_prefixes = [w.split(".")[0] for w in weight_names] + # .bin, .safetensors, ... + weight_suffixs = [w.split(".")[-1] for w in weight_names] + # -00001-of-00002 + transformers_index_format = r"\d{5}-of-\d{5}" + + if variant is not None: + # `diffusion_pytorch_model.fp16.bin` as well as `model.fp16-00001-of-00002.safetensors` + variant_file_re = re.compile( + rf"({'|'.join(weight_prefixes)})\.({variant}|{variant}-{transformers_index_format})\.({'|'.join(weight_suffixs)})$" + ) + # `text_encoder/pytorch_model.bin.index.fp16.json` + variant_index_re = re.compile( + rf"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.{variant}\.json$" + ) + + # `diffusion_pytorch_model.bin` as well as `model-00001-of-00002.safetensors` + non_variant_file_re = re.compile( + rf"({'|'.join(weight_prefixes)})(-{transformers_index_format})?\.({'|'.join(weight_suffixs)})$" + ) + # `text_encoder/pytorch_model.bin.index.json` + non_variant_index_re = re.compile(rf"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.json") + + if variant is not None: + variant_weights = {f for f in filenames if variant_file_re.match(f.split("/")[-1]) is not None} + variant_indexes = {f for f in filenames if variant_index_re.match(f.split("/")[-1]) is not None} + variant_filenames = variant_weights | variant_indexes + else: + variant_filenames = set() + + non_variant_weights = {f for f in filenames if non_variant_file_re.match(f.split("/")[-1]) is not None} + non_variant_indexes = {f for f in filenames if non_variant_index_re.match(f.split("/")[-1]) is not None} + non_variant_filenames = non_variant_weights | non_variant_indexes + + # all variant filenames will be used by default + usable_filenames = set(variant_filenames) + + def convert_to_variant(filename): + if "index" in filename: + variant_filename = filename.replace("index", f"index.{variant}") + elif re.compile(f"^(.*?){transformers_index_format}").match(filename) is not None: + variant_filename = f"{filename.split('-')[0]}.{variant}-{'-'.join(filename.split('-')[1:])}" + else: + variant_filename = f"{filename.split('.')[0]}.{variant}.{filename.split('.')[1]}" + return variant_filename + + for f in non_variant_filenames: + variant_filename = convert_to_variant(f) + if variant_filename not in usable_filenames: + usable_filenames.add(f) + + return usable_filenames, variant_filenames + + +def warn_deprecated_model_variant(pretrained_model_name_or_path, use_auth_token, variant, revision, model_filenames): + info = model_info( + pretrained_model_name_or_path, + use_auth_token=use_auth_token, + revision=None, + ) + filenames = {sibling.rfilename for sibling in info.siblings} + comp_model_filenames, _ = variant_compatible_siblings(filenames, variant=revision) + comp_model_filenames = [".".join(f.split(".")[:1] + f.split(".")[2:]) for f in comp_model_filenames] + + if set(comp_model_filenames) == set(model_filenames): + warnings.warn( + f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'` even though you can load it via `variant=`{revision}`. Loading model variants via `revision='{revision}'` is deprecated and will be removed in diffusers v1. Please use `variant='{revision}'` instead.", + FutureWarning, + ) + else: + warnings.warn( + f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant='{revision}'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have the required variant filenames in the 'main' branch. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title '{pretrained_model_name_or_path} is missing {revision} files' so that the correct variant file can be added.", + FutureWarning, + ) + + +def maybe_raise_or_warn( + library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module +): + """Simple helper method to raise or warn in case incorrect module has been passed""" + if not is_pipeline_module: + library = importlib.import_module(library_name) + class_obj = getattr(library, class_name) + class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} + + expected_class_obj = None + for class_name, class_candidate in class_candidates.items(): + if class_candidate is not None and issubclass(class_obj, class_candidate): + expected_class_obj = class_candidate + + # Dynamo wraps the original model in a private class. + # I didn't find a public API to get the original class. + sub_model = passed_class_obj[name] + model_cls = sub_model.__class__ + if is_compiled_module(sub_model): + model_cls = sub_model._orig_mod.__class__ + + if not issubclass(model_cls, expected_class_obj): + raise ValueError( + f"{passed_class_obj[name]} is of type: {model_cls}, but should be" f" {expected_class_obj}" + ) + else: + logger.warning( + f"You have passed a non-standard module {passed_class_obj[name]}. We cannot verify whether it" + " has the correct type" + ) + + +def get_class_obj_and_candidates(library_name, class_name, importable_classes, pipelines, is_pipeline_module): + """Simple helper method to retrieve class object of module as well as potential parent class objects""" + if is_pipeline_module: + pipeline_module = getattr(pipelines, library_name) + + class_obj = getattr(pipeline_module, class_name) + class_candidates = {c: class_obj for c in importable_classes.keys()} + else: + # else we just import it from the library. + if class_name == 'UNet2DConditionModel': + library_name = "ixformer.contrib.DeepCache.sd.unet_2d_condition" + + + library = importlib.import_module(library_name) + class_obj = getattr(library, class_name) + class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} + + return class_obj, class_candidates + + +def _get_pipeline_class( + class_obj, config, load_connected_pipeline=False, custom_pipeline=None, cache_dir=None, revision=None +): + if custom_pipeline is not None: + if custom_pipeline.endswith(".py"): + path = Path(custom_pipeline) + # decompose into folder & file + file_name = path.name + custom_pipeline = path.parent.absolute() + else: + file_name = CUSTOM_PIPELINE_FILE_NAME + + return get_class_from_dynamic_module( + custom_pipeline, module_file=file_name, cache_dir=cache_dir, revision=revision + ) + + if class_obj != DiffusionPipeline: + return class_obj + + diffusers_module = importlib.import_module(class_obj.__module__.split(".")[0]) + class_name = config["_class_name"] + + if class_name.startswith("Flax"): + class_name = class_name[4:] + + pipeline_cls = getattr(diffusers_module, class_name) + + if load_connected_pipeline: + from .auto_pipeline import _get_connected_pipeline + + connected_pipeline_cls = _get_connected_pipeline(pipeline_cls) + if connected_pipeline_cls is not None: + logger.info( + f"Loading connected pipeline {connected_pipeline_cls.__name__} instead of {pipeline_cls.__name__} as specified via `load_connected_pipeline=True`" + ) + else: + logger.info(f"{pipeline_cls.__name__} has no connected pipeline class. Loading {pipeline_cls.__name__}.") + + pipeline_cls = connected_pipeline_cls or pipeline_cls + + return pipeline_cls + + +def load_sub_model( + library_name: str, + class_name: str, + importable_classes: List[Any], + pipelines: Any, + is_pipeline_module: bool, + pipeline_class: Any, + torch_dtype: torch.dtype, + provider: Any, + sess_options: Any, + device_map: Optional[Union[Dict[str, torch.device], str]], + max_memory: Optional[Dict[Union[int, str], Union[int, str]]], + offload_folder: Optional[Union[str, os.PathLike]], + offload_state_dict: bool, + model_variants: Dict[str, str], + name: str, + from_flax: bool, + variant: str, + low_cpu_mem_usage: bool, + cached_folder: Union[str, os.PathLike], +): + """Helper method to load the module `name` from `library_name` and `class_name`""" + # retrieve class candidates + class_obj, class_candidates = get_class_obj_and_candidates( + library_name, class_name, importable_classes, pipelines, is_pipeline_module + ) + + load_method_name = None + # retrive load method name + for class_name, class_candidate in class_candidates.items(): + if class_candidate is not None and issubclass(class_obj, class_candidate): + load_method_name = importable_classes[class_name][1] + + # if load method name is None, then we have a dummy module -> raise Error + if load_method_name is None: + none_module = class_obj.__module__ + is_dummy_path = none_module.startswith(DUMMY_MODULES_FOLDER) or none_module.startswith( + TRANSFORMERS_DUMMY_MODULES_FOLDER + ) + if is_dummy_path and "dummy" in none_module: + # call class_obj for nice error message of missing requirements + class_obj() + + raise ValueError( + f"The component {class_obj} of {pipeline_class} cannot be loaded as it does not seem to have" + f" any of the loading methods defined in {ALL_IMPORTABLE_CLASSES}." + ) + + load_method = getattr(class_obj, load_method_name) + + # add kwargs to loading method + loading_kwargs = {} + if issubclass(class_obj, torch.nn.Module): + loading_kwargs["torch_dtype"] = torch_dtype + if issubclass(class_obj, diffusers.OnnxRuntimeModel): + loading_kwargs["provider"] = provider + loading_kwargs["sess_options"] = sess_options + + is_diffusers_model = issubclass(class_obj, diffusers.ModelMixin) + + if is_transformers_available(): + transformers_version = version.parse(version.parse(transformers.__version__).base_version) + else: + transformers_version = "N/A" + + is_transformers_model = ( + is_transformers_available() + and issubclass(class_obj, PreTrainedModel) + and transformers_version >= version.parse("4.20.0") + ) + + # When loading a transformers model, if the device_map is None, the weights will be initialized as opposed to diffusers. + # To make default loading faster we set the `low_cpu_mem_usage=low_cpu_mem_usage` flag which is `True` by default. + # This makes sure that the weights won't be initialized which significantly speeds up loading. + if is_diffusers_model or is_transformers_model: + loading_kwargs["device_map"] = device_map + loading_kwargs["max_memory"] = max_memory + loading_kwargs["offload_folder"] = offload_folder + loading_kwargs["offload_state_dict"] = offload_state_dict + loading_kwargs["variant"] = model_variants.pop(name, None) + if from_flax: + loading_kwargs["from_flax"] = True + + # the following can be deleted once the minimum required `transformers` version + # is higher than 4.27 + if ( + is_transformers_model + and loading_kwargs["variant"] is not None + and transformers_version < version.parse("4.27.0") + ): + raise ImportError( + f"When passing `variant='{variant}'`, please make sure to upgrade your `transformers` version to at least 4.27.0.dev0" + ) + elif is_transformers_model and loading_kwargs["variant"] is None: + loading_kwargs.pop("variant") + + # if `from_flax` and model is transformer model, can currently not load with `low_cpu_mem_usage` + if not (from_flax and is_transformers_model): + loading_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage + else: + loading_kwargs["low_cpu_mem_usage"] = False + + # check if the module is in a subdirectory + if os.path.isdir(os.path.join(cached_folder, name)): + loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs) + else: + # else load from the root directory + loaded_sub_model = load_method(cached_folder, **loading_kwargs) + + return loaded_sub_model + + +class DiffusionPipeline(ConfigMixin, PushToHubMixin): + r""" + Base class for all pipelines. + + [`DiffusionPipeline`] stores all components (models, schedulers, and processors) for diffusion pipelines and + provides methods for loading, downloading and saving models. It also includes methods to: + + - move all PyTorch modules to the device of your choice + - enable/disable the progress bar for the denoising iteration + + Class attributes: + + - **config_name** (`str`) -- The configuration filename that stores the class and module names of all the + diffusion pipeline's components. + - **_optional_components** (`List[str]`) -- List of all optional components that don't have to be passed to the + pipeline to function (should be overridden by subclasses). + """ + config_name = "model_index.json" + model_cpu_offload_seq = None + _optional_components = [] + _exclude_from_cpu_offload = [] + _load_connected_pipes = False + _is_onnx = False + + def register_modules(self, **kwargs): + # import it here to avoid circular import + from diffusers import pipelines + + for name, module in kwargs.items(): + # retrieve library + if module is None: + register_dict = {name: (None, None)} + else: + # register the config from the original module, not the dynamo compiled one + if is_compiled_module(module): + not_compiled_module = module._orig_mod + else: + not_compiled_module = module + + library = not_compiled_module.__module__.split(".")[0] + + # check if the module is a pipeline module + module_path_items = not_compiled_module.__module__.split(".") + pipeline_dir = module_path_items[-2] if len(module_path_items) > 2 else None + + path = not_compiled_module.__module__.split(".") + is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir) + + # if library is not in LOADABLE_CLASSES, then it is a custom module. + # Or if it's a pipeline module, then the module is inside the pipeline + # folder so we set the library to module name. + if is_pipeline_module: + library = pipeline_dir + elif library not in LOADABLE_CLASSES: + library = not_compiled_module.__module__ + + # retrieve class_name + class_name = not_compiled_module.__class__.__name__ + + register_dict = {name: (library, class_name)} + + # save model index config + self.register_to_config(**register_dict) + + # set models + setattr(self, name, module) + + def __setattr__(self, name: str, value: Any): + if name in self.__dict__ and hasattr(self.config, name): + # We need to overwrite the config if name exists in config + if isinstance(getattr(self.config, name), (tuple, list)): + if value is not None and self.config[name][0] is not None: + class_library_tuple = (value.__module__.split(".")[0], value.__class__.__name__) + else: + class_library_tuple = (None, None) + + self.register_to_config(**{name: class_library_tuple}) + else: + self.register_to_config(**{name: value}) + + super().__setattr__(name, value) + + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + safe_serialization: bool = True, + variant: Optional[str] = None, + push_to_hub: bool = False, + **kwargs, + ): + """ + Save all saveable variables of the pipeline to a directory. A pipeline variable can be saved and loaded if its + class implements both a save and loading method. The pipeline is easily reloaded using the + [`~DiffusionPipeline.from_pretrained`] class method. + + Arguments: + save_directory (`str` or `os.PathLike`): + Directory to save a pipeline to. Will be created if it doesn't exist. + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. + variant (`str`, *optional*): + If specified, weights are saved in the format `pytorch_model..bin`. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + model_index_dict = dict(self.config) + model_index_dict.pop("_class_name", None) + model_index_dict.pop("_diffusers_version", None) + model_index_dict.pop("_module", None) + model_index_dict.pop("_name_or_path", None) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + private = kwargs.pop("private", False) + create_pr = kwargs.pop("create_pr", False) + token = kwargs.pop("token", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id + + expected_modules, optional_kwargs = self._get_signature_keys(self) + + def is_saveable_module(name, value): + if name not in expected_modules: + return False + if name in self._optional_components and value[0] is None: + return False + return True + + model_index_dict = {k: v for k, v in model_index_dict.items() if is_saveable_module(k, v)} + for pipeline_component_name in model_index_dict.keys(): + sub_model = getattr(self, pipeline_component_name) + model_cls = sub_model.__class__ + + # Dynamo wraps the original model in a private class. + # I didn't find a public API to get the original class. + if is_compiled_module(sub_model): + sub_model = sub_model._orig_mod + model_cls = sub_model.__class__ + + save_method_name = None + # search for the model's base class in LOADABLE_CLASSES + for library_name, library_classes in LOADABLE_CLASSES.items(): + if library_name in sys.modules: + library = importlib.import_module(library_name) + else: + logger.info( + f"{library_name} is not installed. Cannot save {pipeline_component_name} as {library_classes} from {library_name}" + ) + + for base_class, save_load_methods in library_classes.items(): + class_candidate = getattr(library, base_class, None) + if class_candidate is not None and issubclass(model_cls, class_candidate): + # if we found a suitable base class in LOADABLE_CLASSES then grab its save method + save_method_name = save_load_methods[0] + break + if save_method_name is not None: + break + + if save_method_name is None: + logger.warn(f"self.{pipeline_component_name}={sub_model} of type {type(sub_model)} cannot be saved.") + # make sure that unsaveable components are not tried to be loaded afterward + self.register_to_config(**{pipeline_component_name: (None, None)}) + continue + + save_method = getattr(sub_model, save_method_name) + + # Call the save method with the argument safe_serialization only if it's supported + save_method_signature = inspect.signature(save_method) + save_method_accept_safe = "safe_serialization" in save_method_signature.parameters + save_method_accept_variant = "variant" in save_method_signature.parameters + + save_kwargs = {} + if save_method_accept_safe: + save_kwargs["safe_serialization"] = safe_serialization + if save_method_accept_variant: + save_kwargs["variant"] = variant + + save_method(os.path.join(save_directory, pipeline_component_name), **save_kwargs) + + # finally save the config + self.save_config(save_directory) + + if push_to_hub: + self._upload_folder( + save_directory, + repo_id, + token=token, + commit_message=commit_message, + create_pr=create_pr, + ) + + def to( + self, + torch_device: Optional[Union[str, torch.device]] = None, + torch_dtype: Optional[torch.dtype] = None, + silence_dtype_warnings: bool = False, + ): + if torch_device is None and torch_dtype is None: + return self + + # throw warning if pipeline is in "offloaded"-mode but user tries to manually set to GPU. + def module_is_sequentially_offloaded(module): + if not is_accelerate_available() or is_accelerate_version("<", "0.14.0"): + return False + + return hasattr(module, "_hf_hook") and not isinstance( + module._hf_hook, (accelerate.hooks.CpuOffload, accelerate.hooks.AlignDevicesHook) + ) + + def module_is_offloaded(module): + if not is_accelerate_available() or is_accelerate_version("<", "0.17.0.dev0"): + return False + + return hasattr(module, "_hf_hook") and isinstance(module._hf_hook, accelerate.hooks.CpuOffload) + + # .to("cuda") would raise an error if the pipeline is sequentially offloaded, so we raise our own to make it clearer + pipeline_is_sequentially_offloaded = any( + module_is_sequentially_offloaded(module) for _, module in self.components.items() + ) + if pipeline_is_sequentially_offloaded and torch_device and torch.device(torch_device).type == "cuda": + raise ValueError( + "It seems like you have activated sequential model offloading by calling `enable_sequential_cpu_offload`, but are now attempting to move the pipeline to GPU. This is not compatible with offloading. Please, move your pipeline `.to('cpu')` or consider removing the move altogether if you use sequential offloading." + ) + + # Display a warning in this case (the operation succeeds but the benefits are lost) + pipeline_is_offloaded = any(module_is_offloaded(module) for _, module in self.components.items()) + if pipeline_is_offloaded and torch_device and torch.device(torch_device).type == "cuda": + logger.warning( + f"It seems like you have activated model offloading by calling `enable_model_cpu_offload`, but are now manually moving the pipeline to GPU. It is strongly recommended against doing so as memory gains from offloading are likely to be lost. Offloading automatically takes care of moving the individual components {', '.join(self.components.keys())} to GPU when needed. To make sure offloading works as expected, you should consider moving the pipeline back to CPU: `pipeline.to('cpu')` or removing the move altogether if you use offloading." + ) + + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + is_offloaded = pipeline_is_offloaded or pipeline_is_sequentially_offloaded + for module in modules: + is_loaded_in_8bit = hasattr(module, "is_loaded_in_8bit") and module.is_loaded_in_8bit + + if is_loaded_in_8bit and torch_dtype is not None: + logger.warning( + f"The module '{module.__class__.__name__}' has been loaded in 8bit and conversion to {torch_dtype} is not yet supported. Module is still in 8bit precision." + ) + + if is_loaded_in_8bit and torch_device is not None: + logger.warning( + f"The module '{module.__class__.__name__}' has been loaded in 8bit and moving it to {torch_dtype} via `.to()` is not yet supported. Module is still on {module.device}." + ) + else: + module.to(torch_device, torch_dtype) + + if ( + module.dtype == torch.float16 + and str(torch_device) in ["cpu"] + and not silence_dtype_warnings + and not is_offloaded + ): + logger.warning( + "Pipelines loaded with `torch_dtype=torch.float16` cannot run with `cpu` device. It" + " is not recommended to move them to `cpu` as running them will fail. Please make" + " sure to use an accelerator to run the pipeline in inference, due to the lack of" + " support for`float16` operations on this device in PyTorch. Please, remove the" + " `torch_dtype=torch.float16` argument, or use another device for inference." + ) + return self + + @property + def device(self) -> torch.device: + r""" + Returns: + `torch.device`: The torch device on which the pipeline is located. + """ + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + for module in modules: + return module.device + + return torch.device("cpu") + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): + r""" + Instantiate a PyTorch diffusion pipeline from pretrained pipeline weights. + + The pipeline is set in evaluation mode (`model.eval()`) by default. + + If you get the error message below, you need to finetune the weights for your downstream task: + + ``` + Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: + - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated + You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. + ``` + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): + Can be either: + + - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline + hosted on the Hub. + - A path to a *directory* (for example `./my_pipeline_directory/`) containing pipeline weights + saved using + [`~DiffusionPipeline.save_pretrained`]. + torch_dtype (`str` or `torch.dtype`, *optional*): + Override the default `torch.dtype` and load the model with another dtype. If "auto" is passed, the + dtype is automatically derived from the model's weights. + custom_pipeline (`str`, *optional*): + + + + 🧪 This is an experimental feature and may change in the future. + + + + Can be either: + + - A string, the *repo id* (for example `hf-internal-testing/diffusers-dummy-pipeline`) of a custom + pipeline hosted on the Hub. The repository must contain a file called pipeline.py that defines + the custom pipeline. + - A string, the *file name* of a community pipeline hosted on GitHub under + [Community](https://github.com/huggingface/diffusers/tree/main/examples/community). Valid file + names must match the file name and not the pipeline script (`clip_guided_stable_diffusion` + instead of `clip_guided_stable_diffusion.py`). Community pipelines are always loaded from the + current main branch of GitHub. + - A path to a directory (`./my_pipeline_directory/`) containing a custom pipeline. The directory + must contain a file called `pipeline.py` that defines the custom pipeline. + + For more information on how to load and create custom pipelines, please have a look at [Loading and + Adding Custom + Pipelines](https://huggingface.co/docs/diffusers/using-diffusers/custom_pipeline_overview) + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory where a downloaded pretrained model configuration is cached if the standard cache + is not used. + resume_download (`bool`, *optional*, defaults to `False`): + Whether or not to resume downloading the model weights and configuration files. If set to `False`, any + incompletely downloaded files are deleted. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info(`bool`, *optional*, defaults to `False`): + Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files or not. If set to `True`, the model + won't be downloaded from the Hub. + use_auth_token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from + `diffusers-cli login` (stored in `~/.huggingface`) is used. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier + allowed by Git. + custom_revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id similar to + `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a + custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. + mirror (`str`, *optional*): + Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not + guarantee the timeliness or safety of the source, and you should refer to the mirror site for more + information. + device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + A map that specifies where each submodule should go. It doesn’t need to be defined for each + parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the + same device. + + Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For + more information about each option see [designing a device + map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). + max_memory (`Dict`, *optional*): + A dictionary device identifier for the maximum memory. Will default to the maximum memory available for + each GPU and the available CPU RAM if unset. + offload_folder (`str` or `os.PathLike`, *optional*): + The path to offload weights if device_map contains the value `"disk"`. + offload_state_dict (`bool`, *optional*): + If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if + the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` + when there is some disk offload. + low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): + Speed up model loading only loading the pretrained weights and not initializing the weights. This also + tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. + Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this + argument to `True` will raise an error. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `None`, the safetensors weights are downloaded if they're available **and** if the + safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors + weights. If set to `False`, safetensors weights are not loaded. + use_onnx (`bool`, *optional*, defaults to `None`): + If set to `True`, ONNX weights will always be downloaded if present. If set to `False`, ONNX weights + will never be downloaded. By default `use_onnx` defaults to the `_is_onnx` class attribute which is + `False` for non-ONNX pipelines and `True` for ONNX pipelines. ONNX weights include both files ending + with `.onnx` and `.pb`. + kwargs (remaining dictionary of keyword arguments, *optional*): + Can be used to overwrite load and saveable variables (the pipeline components of the specific pipeline + class). The overwritten components are passed directly to the pipelines `__init__` method. See example + below for more information. + variant (`str`, *optional*): + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when + loading `from_flax`. + + + + To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with + `huggingface-cli login`. + + + + Examples: + + ```py + >>> from diffusers import DiffusionPipeline + + >>> # Download pipeline from huggingface.co and cache. + >>> pipeline = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256") + + >>> # Download pipeline that requires an authorization token + >>> # For more information on access tokens, please refer to this section + >>> # of the documentation](https://huggingface.co/docs/hub/security-tokens) + >>> pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5") + + >>> # Use a different scheduler + >>> from diffusers import LMSDiscreteScheduler + + >>> scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config) + >>> pipeline.scheduler = scheduler + ``` + """ + cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE) + resume_download = kwargs.pop("resume_download", False) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + from_flax = kwargs.pop("from_flax", False) + torch_dtype = kwargs.pop("torch_dtype", None) + custom_pipeline = kwargs.pop("custom_pipeline", None) + custom_revision = kwargs.pop("custom_revision", None) + provider = kwargs.pop("provider", None) + sess_options = kwargs.pop("sess_options", None) + device_map = kwargs.pop("device_map", None) + max_memory = kwargs.pop("max_memory", None) + offload_folder = kwargs.pop("offload_folder", None) + offload_state_dict = kwargs.pop("offload_state_dict", False) + low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None) + use_onnx = kwargs.pop("use_onnx", None) + load_connected_pipeline = kwargs.pop("load_connected_pipeline", False) + + # 1. Download the checkpoints and configs + # use snapshot download here to get it working from from_pretrained + if not os.path.isdir(pretrained_model_name_or_path): + cached_folder = cls.download( + pretrained_model_name_or_path, + cache_dir=cache_dir, + resume_download=resume_download, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + from_flax=from_flax, + use_safetensors=use_safetensors, + use_onnx=use_onnx, + custom_pipeline=custom_pipeline, + custom_revision=custom_revision, + variant=variant, + load_connected_pipeline=load_connected_pipeline, + **kwargs, + ) + else: + cached_folder = pretrained_model_name_or_path + + config_dict = cls.load_config(cached_folder) + + # pop out "_ignore_files" as it is only needed for download + config_dict.pop("_ignore_files", None) + + # 2. Define which model components should load variants + # We retrieve the information by matching whether variant + # model checkpoints exist in the subfolders + model_variants = {} + if variant is not None: + for folder in os.listdir(cached_folder): + folder_path = os.path.join(cached_folder, folder) + is_folder = os.path.isdir(folder_path) and folder in config_dict + variant_exists = is_folder and any( + p.split(".")[1].startswith(variant) for p in os.listdir(folder_path) + ) + if variant_exists: + model_variants[folder] = variant + + # 3. Load the pipeline class, if using custom module then load it from the hub + # if we load from explicit class, let's use it + pipeline_class = _get_pipeline_class( + cls, + config_dict, + load_connected_pipeline=load_connected_pipeline, + custom_pipeline=custom_pipeline, + cache_dir=cache_dir, + revision=custom_revision, + ) + + # DEPRECATED: To be removed in 1.0.0 + if pipeline_class.__name__ == "StableDiffusionInpaintPipeline" and version.parse( + version.parse(config_dict["_diffusers_version"]).base_version + ) <= version.parse("0.5.1"): + from diffusers import StableDiffusionInpaintPipeline, StableDiffusionInpaintPipelineLegacy + + pipeline_class = StableDiffusionInpaintPipelineLegacy + + deprecation_message = ( + "You are using a legacy checkpoint for inpainting with Stable Diffusion, therefore we are loading the" + f" {StableDiffusionInpaintPipelineLegacy} class instead of {StableDiffusionInpaintPipeline}. For" + " better inpainting results, we strongly suggest using Stable Diffusion's official inpainting" + " checkpoint: https://huggingface.co/runwayml/stable-diffusion-inpainting instead or adapting your" + f" checkpoint {pretrained_model_name_or_path} to the format of" + " https://huggingface.co/runwayml/stable-diffusion-inpainting. Note that we do not actively maintain" + " the {StableDiffusionInpaintPipelineLegacy} class and will likely remove it in version 1.0.0." + ) + deprecate("StableDiffusionInpaintPipelineLegacy", "1.0.0", deprecation_message, standard_warn=False) + + # 4. Define expected modules given pipeline signature + # and define non-None initialized modules (=`init_kwargs`) + + # some modules can be passed directly to the init + # in this case they are already instantiated in `kwargs` + # extract them here + expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class) + passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} + passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} + + init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs) + + # define init kwargs and make sure that optional component modules are filtered out + init_kwargs = { + k: init_dict.pop(k) + for k in optional_kwargs + if k in init_dict and k not in pipeline_class._optional_components + } + init_kwargs = {**init_kwargs, **passed_pipe_kwargs} + + # remove `null` components + def load_module(name, value): + if value[0] is None: + return False + if name in passed_class_obj and passed_class_obj[name] is None: + return False + return True + + init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)} + + # Special case: safety_checker must be loaded separately when using `from_flax` + if from_flax and "safety_checker" in init_dict and "safety_checker" not in passed_class_obj: + raise NotImplementedError( + "The safety checker cannot be automatically loaded when loading weights `from_flax`." + " Please, pass `safety_checker=None` to `from_pretrained`, and load the safety checker" + " separately if you need it." + ) + + # 5. Throw nice warnings / errors for fast accelerate loading + if len(unused_kwargs) > 0: + logger.warning( + f"Keyword arguments {unused_kwargs} are not expected by {pipeline_class.__name__} and will be ignored." + ) + + if low_cpu_mem_usage and not is_accelerate_available(): + low_cpu_mem_usage = False + logger.warning( + "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" + " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" + " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" + " install accelerate\n```\n." + ) + + if device_map is not None and not is_torch_version(">=", "1.9.0"): + raise NotImplementedError( + "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set" + " `device_map=None`." + ) + + if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): + raise NotImplementedError( + "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" + " `low_cpu_mem_usage=False`." + ) + + if low_cpu_mem_usage is False and device_map is not None: + raise ValueError( + f"You cannot set `low_cpu_mem_usage` to False while using device_map={device_map} for loading and" + " dispatching. Please make sure to set `low_cpu_mem_usage=True`." + ) + + # import it here to avoid circular import + from diffusers import pipelines + + # 6. Load each module in the pipeline + for name, (library_name, class_name) in tqdm(init_dict.items(), desc="Loading pipeline components..."): + # 6.1 - now that JAX/Flax is an official framework of the library, we might load from Flax names + if class_name.startswith("Flax"): + class_name = class_name[4:] + + # 6.2 Define all importable classes + is_pipeline_module = hasattr(pipelines, library_name) + importable_classes = ALL_IMPORTABLE_CLASSES + loaded_sub_model = None + + # 6.3 Use passed sub model or load class_name from library_name + if name in passed_class_obj: + # if the model is in a pipeline module, then we load it from the pipeline + # check that passed_class_obj has correct parent class + maybe_raise_or_warn( + library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module + ) + + loaded_sub_model = passed_class_obj[name] + else: + # load sub model + loaded_sub_model = load_sub_model( + library_name=library_name, + class_name=class_name, + importable_classes=importable_classes, + pipelines=pipelines, + is_pipeline_module=is_pipeline_module, + pipeline_class=pipeline_class, + torch_dtype=torch_dtype, + provider=provider, + sess_options=sess_options, + device_map=device_map, + max_memory=max_memory, + offload_folder=offload_folder, + offload_state_dict=offload_state_dict, + model_variants=model_variants, + name=name, + from_flax=from_flax, + variant=variant, + low_cpu_mem_usage=low_cpu_mem_usage, + cached_folder=cached_folder, + ) + logger.info( + f"Loaded {name} as {class_name} from `{name}` subfolder of {pretrained_model_name_or_path}." + ) + + init_kwargs[name] = loaded_sub_model # UNet(...), # DiffusionSchedule(...) + + if pipeline_class._load_connected_pipes and os.path.isfile(os.path.join(cached_folder, "README.md")): + modelcard = ModelCard.load(os.path.join(cached_folder, "README.md")) + connected_pipes = {prefix: getattr(modelcard.data, prefix, [None])[0] for prefix in CONNECTED_PIPES_KEYS} + load_kwargs = { + "cache_dir": cache_dir, + "resume_download": resume_download, + "force_download": force_download, + "proxies": proxies, + "local_files_only": local_files_only, + "use_auth_token": use_auth_token, + "revision": revision, + "torch_dtype": torch_dtype, + "custom_pipeline": custom_pipeline, + "custom_revision": custom_revision, + "provider": provider, + "sess_options": sess_options, + "device_map": device_map, + "max_memory": max_memory, + "offload_folder": offload_folder, + "offload_state_dict": offload_state_dict, + "low_cpu_mem_usage": low_cpu_mem_usage, + "variant": variant, + "use_safetensors": use_safetensors, + } + + def get_connected_passed_kwargs(prefix): + connected_passed_class_obj = { + k.replace(f"{prefix}_", ""): w for k, w in passed_class_obj.items() if k.split("_")[0] == prefix + } + connected_passed_pipe_kwargs = { + k.replace(f"{prefix}_", ""): w for k, w in passed_pipe_kwargs.items() if k.split("_")[0] == prefix + } + + connected_passed_kwargs = {**connected_passed_class_obj, **connected_passed_pipe_kwargs} + return connected_passed_kwargs + + connected_pipes = { + prefix: DiffusionPipeline.from_pretrained( + repo_id, **load_kwargs.copy(), **get_connected_passed_kwargs(prefix) + ) + for prefix, repo_id in connected_pipes.items() + if repo_id is not None + } + + for prefix, connected_pipe in connected_pipes.items(): + # add connected pipes to `init_kwargs` with _, e.g. "prior_text_encoder" + init_kwargs.update( + {"_".join([prefix, name]): component for name, component in connected_pipe.components.items()} + ) + + # 7. Potentially add passed objects if expected + missing_modules = set(expected_modules) - set(init_kwargs.keys()) + passed_modules = list(passed_class_obj.keys()) + optional_modules = pipeline_class._optional_components + if len(missing_modules) > 0 and missing_modules <= set(passed_modules + optional_modules): + for module in missing_modules: + init_kwargs[module] = passed_class_obj.get(module, None) + elif len(missing_modules) > 0: + passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys())) - optional_kwargs + raise ValueError( + f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed." + ) + + # 8. Instantiate the pipeline + model = pipeline_class(**init_kwargs) + + # 9. Save where the model was instantiated from + model.register_to_config(_name_or_path=pretrained_model_name_or_path) + return model + + @property + def name_or_path(self) -> str: + return getattr(self.config, "_name_or_path", None) + + @property + def _execution_device(self): + r""" + Returns the device on which the pipeline's models will be executed. After calling + [`~DiffusionPipeline.enable_sequential_cpu_offload`] the execution device can only be inferred from + Accelerate's module hooks. + """ + for name, model in self.components.items(): + if not isinstance(model, torch.nn.Module) or name in self._exclude_from_cpu_offload: + continue + + if not hasattr(model, "_hf_hook"): + return self.device + for module in model.modules(): + if ( + hasattr(module, "_hf_hook") + and hasattr(module._hf_hook, "execution_device") + and module._hf_hook.execution_device is not None + ): + return torch.device(module._hf_hook.execution_device) + return self.device + + def enable_model_cpu_offload(self, gpu_id: int = 0, device: Union[torch.device, str] = "cuda"): + r""" + Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared + to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward` + method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with + `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`. + """ + if self.model_cpu_offload_seq is None: + raise ValueError( + "Model CPU offload cannot be enabled because no `model_cpu_offload_seq` class attribute is set." + ) + + if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"): + from accelerate import cpu_offload_with_hook + else: + raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.") + + device = torch.device(f"cuda:{gpu_id}") + + if self.device.type != "cpu": + self.to("cpu", silence_dtype_warnings=True) + device_mod = getattr(torch, self.device.type, None) + if hasattr(device_mod, "empty_cache") and device_mod.is_available(): + device_mod.empty_cache() # otherwise we don't see the memory savings (but they probably exist) + + all_model_components = {k: v for k, v in self.components.items() if isinstance(v, torch.nn.Module)} + + self._all_hooks = [] + hook = None + for model_str in self.model_cpu_offload_seq.split("->"): + model = all_model_components.pop(model_str, None) + if not isinstance(model, torch.nn.Module): + continue + + _, hook = cpu_offload_with_hook(model, device, prev_module_hook=hook) + self._all_hooks.append(hook) + + # CPU offload models that are not in the seq chain unless they are explicitly excluded + # these models will stay on CPU until maybe_free_model_hooks is called + # some models cannot be in the seq chain because they are iteratively called, such as controlnet + for name, model in all_model_components.items(): + if not isinstance(model, torch.nn.Module): + continue + + if name in self._exclude_from_cpu_offload: + model.to(device) + else: + _, hook = cpu_offload_with_hook(model, device) + self._all_hooks.append(hook) + + def maybe_free_model_hooks(self): + r""" + TODO: Better doc string + """ + if not hasattr(self, "_all_hooks") or len(self._all_hooks) == 0: + # `enable_model_cpu_offload` has not be called, so silently do nothing + return + + for hook in self._all_hooks: + # offload model and remove hook from model + hook.offload() + hook.remove() + + # make sure the model is in the same state as before calling it + self.enable_model_cpu_offload() + + def enable_sequential_cpu_offload(self, gpu_id: int = 0, device: Union[torch.device, str] = "cuda"): + r""" + Offloads all models to CPU using 🤗 Accelerate, significantly reducing memory usage. When called, the state + dicts of all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are saved to CPU + and then moved to `torch.device('meta')` and loaded to GPU only when their specific submodule has its `forward` + method called. Offloading happens on a submodule basis. Memory savings are higher than with + `enable_model_cpu_offload`, but performance is lower. + """ + if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"): + from accelerate import cpu_offload + else: + raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher") + + if device == "cuda": + device = torch.device(f"{device}:{gpu_id}") + + if self.device.type != "cpu": + self.to("cpu", silence_dtype_warnings=True) + device_mod = getattr(torch, self.device.type, None) + if hasattr(device_mod, "empty_cache") and device_mod.is_available(): + device_mod.empty_cache() # otherwise we don't see the memory savings (but they probably exist) + + for name, model in self.components.items(): + if not isinstance(model, torch.nn.Module): + continue + + if name in self._exclude_from_cpu_offload: + model.to(device) + else: + # make sure to offload buffers if not all high level weights + # are of type nn.Module + offload_buffers = len(model._parameters) > 0 + cpu_offload(model, device, offload_buffers=offload_buffers) + + @classmethod + def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: + r""" + Download and cache a PyTorch diffusion pipeline from pretrained pipeline weights. + + Parameters: + pretrained_model_name (`str` or `os.PathLike`, *optional*): + A string, the *repository id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline + hosted on the Hub. + custom_pipeline (`str`, *optional*): + Can be either: + + - A string, the *repository id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained + pipeline hosted on the Hub. The repository must contain a file called `pipeline.py` that defines + the custom pipeline. + + - A string, the *file name* of a community pipeline hosted on GitHub under + [Community](https://github.com/huggingface/diffusers/tree/main/examples/community). Valid file + names must match the file name and not the pipeline script (`clip_guided_stable_diffusion` + instead of `clip_guided_stable_diffusion.py`). Community pipelines are always loaded from the + current `main` branch of GitHub. + + - A path to a *directory* (`./my_pipeline_directory/`) containing a custom pipeline. The directory + must contain a file called `pipeline.py` that defines the custom pipeline. + + + + 🧪 This is an experimental feature and may change in the future. + + + + For more information on how to load and create custom pipelines, take a look at [How to contribute a + community pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/contribute_pipeline). + + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + resume_download (`bool`, *optional*, defaults to `False`): + Whether or not to resume downloading the model weights and configuration files. If set to `False`, any + incompletely downloaded files are deleted. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info(`bool`, *optional*, defaults to `False`): + Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files or not. If set to `True`, the model + won't be downloaded from the Hub. + use_auth_token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from + `diffusers-cli login` (stored in `~/.huggingface`) is used. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier + allowed by Git. + custom_revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id similar to + `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a + custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. + mirror (`str`, *optional*): + Mirror source to resolve accessibility issues if you're downloading a model in China. We do not + guarantee the timeliness or safety of the source, and you should refer to the mirror site for more + information. + variant (`str`, *optional*): + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when + loading `from_flax`. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `None`, the safetensors weights are downloaded if they're available **and** if the + safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors + weights. If set to `False`, safetensors weights are not loaded. + use_onnx (`bool`, *optional*, defaults to `False`): + If set to `True`, ONNX weights will always be downloaded if present. If set to `False`, ONNX weights + will never be downloaded. By default `use_onnx` defaults to the `_is_onnx` class attribute which is + `False` for non-ONNX pipelines and `True` for ONNX pipelines. ONNX weights include both files ending + with `.onnx` and `.pb`. + + Returns: + `os.PathLike`: + A path to the downloaded pipeline. + + + + To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with + `huggingface-cli login`. + + + + """ + cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE) + resume_download = kwargs.pop("resume_download", False) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + from_flax = kwargs.pop("from_flax", False) + custom_pipeline = kwargs.pop("custom_pipeline", None) + custom_revision = kwargs.pop("custom_revision", None) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None) + use_onnx = kwargs.pop("use_onnx", None) + load_connected_pipeline = kwargs.pop("load_connected_pipeline", False) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = True + allow_pickle = True + + allow_patterns = None + ignore_patterns = None + + model_info_call_error: Optional[Exception] = None + if not local_files_only: + try: + info = model_info( + pretrained_model_name, + use_auth_token=use_auth_token, + revision=revision, + ) + except HTTPError as e: + logger.warn(f"Couldn't connect to the Hub: {e}.\nWill try to load from local cache.") + local_files_only = True + model_info_call_error = e # save error to reraise it if model is not cached locally + + if not local_files_only: + config_file = hf_hub_download( + pretrained_model_name, + cls.config_name, + cache_dir=cache_dir, + revision=revision, + proxies=proxies, + force_download=force_download, + resume_download=resume_download, + use_auth_token=use_auth_token, + ) + + config_dict = cls._dict_from_json_file(config_file) + + ignore_filenames = config_dict.pop("_ignore_files", []) + + # retrieve all folder_names that contain relevant files + folder_names = [k for k, v in config_dict.items() if isinstance(v, list)] + + filenames = {sibling.rfilename for sibling in info.siblings} + model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant) + + if len(variant_filenames) == 0 and variant is not None: + deprecation_message = ( + f"You are trying to load the model files of the `variant={variant}`, but no such modeling files are available." + f"The default model files: {model_filenames} will be loaded instead. Make sure to not load from `variant={variant}`" + "if such variant modeling files are not available. Doing so will lead to an error in v0.22.0 as defaulting to non-variant" + "modeling files is deprecated." + ) + deprecate("no variant default", "0.22.0", deprecation_message, standard_warn=False) + + # remove ignored filenames + model_filenames = set(model_filenames) - set(ignore_filenames) + variant_filenames = set(variant_filenames) - set(ignore_filenames) + + # if the whole pipeline is cached we don't have to ping the Hub + if revision in DEPRECATED_REVISION_ARGS and version.parse( + version.parse(__version__).base_version + ) >= version.parse("0.22.0"): + warn_deprecated_model_variant( + pretrained_model_name, use_auth_token, variant, revision, model_filenames + ) + + model_folder_names = {os.path.split(f)[0] for f in model_filenames if os.path.split(f)[0] in folder_names} + + # all filenames compatible with variant will be added + allow_patterns = list(model_filenames) + + # allow all patterns from non-model folders + # this enables downloading schedulers, tokenizers, ... + allow_patterns += [f"{k}/*" for k in folder_names if k not in model_folder_names] + # also allow downloading config.json files with the model + allow_patterns += [os.path.join(k, "config.json") for k in model_folder_names] + + allow_patterns += [ + SCHEDULER_CONFIG_NAME, + CONFIG_NAME, + cls.config_name, + CUSTOM_PIPELINE_FILE_NAME, + ] + + # retrieve passed components that should not be downloaded + pipeline_class = _get_pipeline_class( + cls, + config_dict, + load_connected_pipeline=load_connected_pipeline, + custom_pipeline=custom_pipeline, + cache_dir=cache_dir, + revision=custom_revision, + ) + expected_components, _ = cls._get_signature_keys(pipeline_class) + passed_components = [k for k in expected_components if k in kwargs] + + if ( + use_safetensors + and not allow_pickle + and not is_safetensors_compatible( + model_filenames, variant=variant, passed_components=passed_components + ) + ): + raise EnvironmentError( + f"Could not found the necessary `safetensors` weights in {model_filenames} (variant={variant})" + ) + if from_flax: + ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb"] + elif use_safetensors and is_safetensors_compatible( + model_filenames, variant=variant, passed_components=passed_components + ): + ignore_patterns = ["*.bin", "*.msgpack"] + + use_onnx = use_onnx if use_onnx is not None else pipeline_class._is_onnx + if not use_onnx: + ignore_patterns += ["*.onnx", "*.pb"] + + safetensors_variant_filenames = {f for f in variant_filenames if f.endswith(".safetensors")} + safetensors_model_filenames = {f for f in model_filenames if f.endswith(".safetensors")} + if ( + len(safetensors_variant_filenames) > 0 + and safetensors_model_filenames != safetensors_variant_filenames + ): + logger.warn( + f"\nA mixture of {variant} and non-{variant} filenames will be loaded.\nLoaded {variant} filenames:\n[{', '.join(safetensors_variant_filenames)}]\nLoaded non-{variant} filenames:\n[{', '.join(safetensors_model_filenames - safetensors_variant_filenames)}\nIf this behavior is not expected, please check your folder structure." + ) + else: + ignore_patterns = ["*.safetensors", "*.msgpack"] + + use_onnx = use_onnx if use_onnx is not None else pipeline_class._is_onnx + if not use_onnx: + ignore_patterns += ["*.onnx", "*.pb"] + + bin_variant_filenames = {f for f in variant_filenames if f.endswith(".bin")} + bin_model_filenames = {f for f in model_filenames if f.endswith(".bin")} + if len(bin_variant_filenames) > 0 and bin_model_filenames != bin_variant_filenames: + logger.warn( + f"\nA mixture of {variant} and non-{variant} filenames will be loaded.\nLoaded {variant} filenames:\n[{', '.join(bin_variant_filenames)}]\nLoaded non-{variant} filenames:\n[{', '.join(bin_model_filenames - bin_variant_filenames)}\nIf this behavior is not expected, please check your folder structure." + ) + + # Don't download any objects that are passed + allow_patterns = [ + p for p in allow_patterns if not (len(p.split("/")) == 2 and p.split("/")[0] in passed_components) + ] + + if pipeline_class._load_connected_pipes: + allow_patterns.append("README.md") + + # Don't download index files of forbidden patterns either + ignore_patterns = ignore_patterns + [f"{i}.index.*json" for i in ignore_patterns] + + re_ignore_pattern = [re.compile(fnmatch.translate(p)) for p in ignore_patterns] + re_allow_pattern = [re.compile(fnmatch.translate(p)) for p in allow_patterns] + + expected_files = [f for f in filenames if not any(p.match(f) for p in re_ignore_pattern)] + expected_files = [f for f in expected_files if any(p.match(f) for p in re_allow_pattern)] + + snapshot_folder = Path(config_file).parent + pipeline_is_cached = all((snapshot_folder / f).is_file() for f in expected_files) + + if pipeline_is_cached and not force_download: + # if the pipeline is cached, we can directly return it + # else call snapshot_download + return snapshot_folder + + user_agent = {"pipeline_class": cls.__name__} + if custom_pipeline is not None and not custom_pipeline.endswith(".py"): + user_agent["custom_pipeline"] = custom_pipeline + + # download all allow_patterns - ignore_patterns + try: + cached_folder = snapshot_download( + pretrained_model_name, + cache_dir=cache_dir, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + user_agent=user_agent, + ) + + # retrieve pipeline class from local file + cls_name = cls.load_config(os.path.join(cached_folder, "model_index.json")).get("_class_name", None) + pipeline_class = getattr(diffusers, cls_name, None) + + if pipeline_class is not None and pipeline_class._load_connected_pipes: + modelcard = ModelCard.load(os.path.join(cached_folder, "README.md")) + connected_pipes = sum([getattr(modelcard.data, k, []) for k in CONNECTED_PIPES_KEYS], []) + for connected_pipe_repo_id in connected_pipes: + download_kwargs = { + "cache_dir": cache_dir, + "resume_download": resume_download, + "force_download": force_download, + "proxies": proxies, + "local_files_only": local_files_only, + "use_auth_token": use_auth_token, + "variant": variant, + "use_safetensors": use_safetensors, + } + DiffusionPipeline.download(connected_pipe_repo_id, **download_kwargs) + + return cached_folder + + except FileNotFoundError: + # Means we tried to load pipeline with `local_files_only=True` but the files have not been found in local cache. + # This can happen in two cases: + # 1. If the user passed `local_files_only=True` => we raise the error directly + # 2. If we forced `local_files_only=True` when `model_info` failed => we raise the initial error + if model_info_call_error is None: + # 1. user passed `local_files_only=True` + raise + else: + # 2. we forced `local_files_only=True` when `model_info` failed + raise EnvironmentError( + f"Cannot load model {pretrained_model_name}: model is not cached locally and an error occured" + " while trying to fetch metadata from the Hub. Please check out the root cause in the stacktrace" + " above." + ) from model_info_call_error + + @staticmethod + def _get_signature_keys(obj): + parameters = inspect.signature(obj.__init__).parameters + required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty} + optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty}) + expected_modules = set(required_parameters.keys()) - {"self"} + return expected_modules, optional_parameters + + @property + def components(self) -> Dict[str, Any]: + r""" + The `self.components` property can be useful to run different pipelines with the same weights and + configurations without reallocating additional memory. + + Returns (`dict`): + A dictionary containing all the modules needed to initialize the pipeline. + + Examples: + + ```py + >>> from diffusers import ( + ... StableDiffusionPipeline, + ... StableDiffusionImg2ImgPipeline, + ... StableDiffusionInpaintPipeline, + ... ) + + >>> text2img = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> img2img = StableDiffusionImg2ImgPipeline(**text2img.components) + >>> inpaint = StableDiffusionInpaintPipeline(**text2img.components) + ``` + """ + expected_modules, optional_parameters = self._get_signature_keys(self) + components = { + k: getattr(self, k) for k in self.config.keys() if not k.startswith("_") and k not in optional_parameters + } + + if set(components.keys()) != expected_modules: + raise ValueError( + f"{self} has been incorrectly initialized or {self.__class__} is incorrectly implemented. Expected" + f" {expected_modules} to be defined, but {components.keys()} are defined." + ) + + return components + + @staticmethod + def numpy_to_pil(images): + """ + Convert a NumPy image or a batch of images to a PIL image. + """ + return numpy_to_pil(images) + + def progress_bar(self, iterable=None, total=None): + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + elif not isinstance(self._progress_bar_config, dict): + raise ValueError( + f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." + ) + + if iterable is not None: + return tqdm(iterable, **self._progress_bar_config) + elif total is not None: + return tqdm(total=total, **self._progress_bar_config) + else: + raise ValueError("Either `total` or `iterable` has to be defined.") + + def set_progress_bar_config(self, **kwargs): + self._progress_bar_config = kwargs + + def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + r""" + Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). When this + option is enabled, you should observe lower GPU memory usage and a potential speed up during inference. Speed + up during training is not guaranteed. + + + + ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes + precedent. + + + + Parameters: + attention_op (`Callable`, *optional*): + Override the default `None` operator for use as `op` argument to the + [`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention) + function of xFormers. + + Examples: + + ```py + >>> import torch + >>> from diffusers import DiffusionPipeline + >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp + + >>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16) + >>> pipe = pipe.to("cuda") + >>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp) + >>> # Workaround for not accepting attention shape using VAE for Flash Attention + >>> pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None) + ``` + """ + self.set_use_memory_efficient_attention_xformers(True, attention_op) + + def disable_xformers_memory_efficient_attention(self): + r""" + Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). + """ + self.set_use_memory_efficient_attention_xformers(False) + + def set_use_memory_efficient_attention_xformers( + self, valid: bool, attention_op: Optional[Callable] = None + ) -> None: + # Recursively walk through all the children. + # Any children which exposes the set_use_memory_efficient_attention_xformers method + # gets the message + def fn_recursive_set_mem_eff(module: torch.nn.Module): + if hasattr(module, "set_use_memory_efficient_attention_xformers"): + module.set_use_memory_efficient_attention_xformers(valid, attention_op) + + for child in module.children(): + fn_recursive_set_mem_eff(child) + + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + for module in modules: + fn_recursive_set_mem_eff(module) + + def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"): + r""" + Enable sliced attention computation. When this option is enabled, the attention module splits the input tensor + in slices to compute attention in several steps. For more than one attention head, the computation is performed + sequentially over each head. This is useful to save some memory in exchange for a small speed decrease. + + + + ⚠️ Don't enable attention slicing if you're already using `scaled_dot_product_attention` (SDPA) from PyTorch + 2.0 or xFormers. These attention computations are already very memory efficient so you won't need to enable + this function. If you enable attention slicing with SDPA or xFormers, it can lead to serious slow downs! + + + + Args: + slice_size (`str` or `int`, *optional*, defaults to `"auto"`): + When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If + `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + + Examples: + + ```py + >>> import torch + >>> from diffusers import StableDiffusionPipeline + + >>> pipe = StableDiffusionPipeline.from_pretrained( + ... "runwayml/stable-diffusion-v1-5", + ... torch_dtype=torch.float16, + ... use_safetensors=True, + ... ) + + >>> prompt = "a photo of an astronaut riding a horse on mars" + >>> pipe.enable_attention_slicing() + >>> image = pipe(prompt).images[0] + ``` + """ + self.set_attention_slice(slice_size) + + def disable_attention_slicing(self): + r""" + Disable sliced attention computation. If `enable_attention_slicing` was previously called, attention is + computed in one step. + """ + # set slice_size = `None` to disable `attention slicing` + self.enable_attention_slicing(None) + + def set_attention_slice(self, slice_size: Optional[int]): + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module) and hasattr(m, "set_attention_slice")] + + for module in modules: + module.set_attention_slice(slice_size) diff --git a/ixformer_sdk/contrib/DeepCache/sd/unet_2d_blocks.py b/ixformer_sdk/contrib/DeepCache/sd/unet_2d_blocks.py new file mode 100644 index 00000000..efb4eb8f --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sd/unet_2d_blocks.py @@ -0,0 +1,3296 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. +from typing import Any, Dict, Optional, Tuple + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from diffusers.utils import is_torch_version, logging +from diffusers.models.activations import get_activation +import diffusers +if diffusers.__version__ >= '0.22.0': + from diffusers.models.normalization import AdaGroupNorm +else: + from diffusers.models.attention import AdaGroupNorm +from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0 +from diffusers.models.dual_transformer_2d import DualTransformer2DModel +from diffusers.models.resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, Upsample2D +from diffusers.models.transformer_2d import Transformer2DModel + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +import time + +def get_down_block( + down_block_type, + num_layers, + in_channels, + out_channels, + temb_channels, + add_downsample, + resnet_eps, + resnet_act_fn, + transformer_layers_per_block=1, + num_attention_heads=None, + resnet_groups=None, + cross_attention_dim=None, + downsample_padding=None, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + resnet_time_scale_shift="default", + attention_type="default", + resnet_skip_time_act=False, + resnet_out_scale_factor=1.0, + cross_attention_norm=None, + attention_head_dim=None, + downsample_type=None, + dropout=0.0, +): + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type + if down_block_type == "DownBlock2D": + return DownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "ResnetDownsampleBlock2D": + return ResnetDownsampleBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + ) + elif down_block_type == "AttnDownBlock2D": + if add_downsample is False: + downsample_type = None + else: + downsample_type = downsample_type or "conv" # default to 'conv' + return AttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + downsample_type=downsample_type, + ) + elif down_block_type == "CrossAttnDownBlock2D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D") + return CrossAttnDownBlock2D( + num_layers=num_layers, + transformer_layers_per_block=transformer_layers_per_block, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + ) + elif down_block_type == "SimpleCrossAttnDownBlock2D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D") + return SimpleCrossAttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif down_block_type == "SkipDownBlock2D": + return SkipDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "AttnSkipDownBlock2D": + return AttnSkipDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "DownEncoderBlock2D": + return DownEncoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "AttnDownEncoderBlock2D": + return AttnDownEncoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "KDownBlock2D": + return KDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + ) + elif down_block_type == "KCrossAttnDownBlock2D": + return KCrossAttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + cross_attention_dim=cross_attention_dim, + attention_head_dim=attention_head_dim, + add_self_attention=True if not add_downsample else False, + ) + raise ValueError(f"{down_block_type} does not exist.") + + +def get_up_block( + up_block_type, + num_layers, + in_channels, + out_channels, + prev_output_channel, + temb_channels, + add_upsample, + resnet_eps, + resnet_act_fn, + transformer_layers_per_block=1, + num_attention_heads=None, + resnet_groups=None, + cross_attention_dim=None, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + resnet_time_scale_shift="default", + attention_type="default", + resnet_skip_time_act=False, + resnet_out_scale_factor=1.0, + cross_attention_norm=None, + attention_head_dim=None, + upsample_type=None, + dropout=0.0, +): + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type + if up_block_type == "UpBlock2D": + return UpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == "ResnetUpsampleBlock2D": + return ResnetUpsampleBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + ) + elif up_block_type == "CrossAttnUpBlock2D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D") + return CrossAttnUpBlock2D( + num_layers=num_layers, + transformer_layers_per_block=transformer_layers_per_block, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + ) + elif up_block_type == "SimpleCrossAttnUpBlock2D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D") + return SimpleCrossAttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif up_block_type == "AttnUpBlock2D": + if add_upsample is False: + upsample_type = None + else: + upsample_type = upsample_type or "conv" # default to 'conv' + + return AttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + upsample_type=upsample_type, + ) + elif up_block_type == "SkipUpBlock2D": + return SkipUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == "AttnSkipUpBlock2D": + return AttnSkipUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == "UpDecoderBlock2D": + return UpDecoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + temb_channels=temb_channels, + ) + elif up_block_type == "AttnUpDecoderBlock2D": + return AttnUpDecoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + temb_channels=temb_channels, + ) + elif up_block_type == "KUpBlock2D": + return KUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + ) + elif up_block_type == "KCrossAttnUpBlock2D": + return KCrossAttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + cross_attention_dim=cross_attention_dim, + attention_head_dim=attention_head_dim, + ) + + raise ValueError(f"{up_block_type} does not exist.") + + +class AutoencoderTinyBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int, act_fn: str): + super().__init__() + act_fn = get_activation(act_fn) + self.conv = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), + act_fn, + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + act_fn, + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + ) + self.skip = ( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) + if in_channels != out_channels + else nn.Identity() + ) + self.fuse = nn.ReLU() + + def forward(self, x): + return self.fuse(self.conv(x) + self.skip(x)) + + +class UNetMidBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + ): + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + self.add_attention = add_attention + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." + ) + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups if resnet_time_scale_shift == "default" else None, + spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + else: + attentions.append(None) + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, hidden_states, temb=None): + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + hidden_states = attn(hidden_states, temb=temb) + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class UNetMidBlock2DCrossAttn(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads=1, + output_scale_factor=1.0, + cross_attention_dim=1280, + dual_cross_attention=False, + use_linear_projection=False, + upcast_attention=False, + attention_type="default", + ): + super().__init__() + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + for _ in range(num_layers): + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + in_channels // num_attention_heads, + in_channels=in_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + in_channels // num_attention_heads, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + else: + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + return hidden_states + + +class UNetMidBlock2DSimpleCrossAttn(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + cross_attention_dim=1280, + skip_time_act=False, + only_cross_attention=False, + cross_attention_norm=None, + ): + super().__init__() + + self.has_cross_attention = True + + self.attention_head_dim = attention_head_dim + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + + self.num_heads = in_channels // self.attention_head_dim + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ] + attentions = [] + + for _ in range(num_layers): + processor = ( + AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() + ) + + attentions.append( + Attention( + query_dim=in_channels, + cross_attention_dim=in_channels, + heads=self.num_heads, + dim_head=self.attention_head_dim, + added_kv_proj_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + bias=True, + upcast_softmax=True, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + processor=processor, + ) + ) + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + lora_scale = cross_attention_kwargs.get("scale", 1.0) + + if attention_mask is None: + # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. + mask = None if encoder_hidden_states is None else encoder_attention_mask + else: + # when attention_mask is defined: we don't even check for encoder_attention_mask. + # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. + # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. + # then we can simplify this whole if/else block to: + # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask + mask = attention_mask + + hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + # attn + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + + # resnet + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + return hidden_states + + +class AttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + downsample_padding=1, + downsample_type="conv", + ): + super().__init__() + resnets = [] + attentions = [] + self.downsample_type = downsample_type + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if downsample_type == "conv": + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + elif downsample_type == "resnet": + self.downsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + down=True, + ) + ] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states, temb=None, upsample_size=None, cross_attention_kwargs=None): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + lora_scale = cross_attention_kwargs.get("scale", 1.0) + + output_states = () + + for resnet, attn in zip(self.resnets, self.attentions): + cross_attention_kwargs.update({"scale": lora_scale}) + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn(hidden_states, **cross_attention_kwargs) + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + if self.downsample_type == "resnet": + hidden_states = downsampler(hidden_states, temb=temb, scale=lora_scale) + else: + hidden_states = downsampler(hidden_states, scale=lora_scale) + + output_states += (hidden_states,) + + return hidden_states, output_states + + +class CrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + downsample_padding=1, + add_downsample=True, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + attention_type="default", + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + exist_block_number=None, + additional_residuals=None, + ): + output_states = () + + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + blocks = list(zip(self.resnets, self.attentions)) + + for i, (resnet, attn) in enumerate(blocks): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + # apply additional residuals to the output of the last pair of resnet and attention blocks + if i == len(blocks) - 1 and additional_residuals is not None: + hidden_states = hidden_states + additional_residuals + + output_states = output_states + (hidden_states,) + if exist_block_number is not None and len(output_states) == exist_block_number + 1: + return hidden_states, output_states + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale=lora_scale) + + output_states = output_states + (hidden_states,) + return hidden_states, output_states + + +class DownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None, scale: float = 1.0): + output_states = () + + i = 0 + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + output_states = output_states + (hidden_states,) + i += 1 + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale=scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class DownEncoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states, scale: float = 1.0): + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None, scale=scale) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale) + + return hidden_states + + +class AttnDownEncoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states, scale: float = 1.0): + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb=None, scale=scale) + cross_attention_kwargs = {"scale": scale} + hidden_states = attn(hidden_states, **cross_attention_kwargs) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale) + + return hidden_states + + +class AttnSkipDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=np.sqrt(2.0), + add_downsample=True, + ): + super().__init__() + self.attentions = nn.ModuleList([]) + self.resnets = nn.ModuleList([]) + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(in_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + self.attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=32, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + if add_downsample: + self.resnet_down = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + down=True, + kernel="fir", + ) + self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)]) + self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1)) + else: + self.resnet_down = None + self.downsamplers = None + self.skip_conv = None + + def forward(self, hidden_states, temb=None, skip_sample=None, scale: float = 1.0): + output_states = () + + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb, scale=scale) + cross_attention_kwargs = {"scale": scale} + hidden_states = attn(hidden_states, **cross_attention_kwargs) + output_states += (hidden_states,) + + if self.downsamplers is not None: + hidden_states = self.resnet_down(hidden_states, temb, scale=scale) + for downsampler in self.downsamplers: + skip_sample = downsampler(skip_sample) + + hidden_states = self.skip_conv(skip_sample) + hidden_states + + output_states += (hidden_states,) + + return hidden_states, output_states, skip_sample + + +class SkipDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_pre_norm: bool = True, + output_scale_factor=np.sqrt(2.0), + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(in_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + if add_downsample: + self.resnet_down = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + down=True, + kernel="fir", + ) + self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)]) + self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1)) + else: + self.resnet_down = None + self.downsamplers = None + self.skip_conv = None + + def forward(self, hidden_states, temb=None, skip_sample=None, scale: float = 1.0): + output_states = () + + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb, scale) + output_states += (hidden_states,) + + if self.downsamplers is not None: + hidden_states = self.resnet_down(hidden_states, temb, scale) + for downsampler in self.downsamplers: + skip_sample = downsampler(skip_sample) + + hidden_states = self.skip_conv(skip_sample) + hidden_states + + output_states += (hidden_states,) + + return hidden_states, output_states, skip_sample + + +class ResnetDownsampleBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + skip_time_act=False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + down=True, + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None, scale: float = 1.0): + output_states = () + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale) + + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, temb, scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class SimpleCrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + add_downsample=True, + skip_time_act=False, + only_cross_attention=False, + cross_attention_norm=None, + ): + super().__init__() + + self.has_cross_attention = True + + resnets = [] + attentions = [] + + self.attention_head_dim = attention_head_dim + self.num_heads = out_channels // self.attention_head_dim + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + processor = ( + AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() + ) + + attentions.append( + Attention( + query_dim=out_channels, + cross_attention_dim=out_channels, + heads=self.num_heads, + dim_head=attention_head_dim, + added_kv_proj_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + bias=True, + upcast_softmax=True, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + processor=processor, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + down=True, + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + output_states = () + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + lora_scale = cross_attention_kwargs.get("scale", 1.0) + + if attention_mask is None: + # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. + mask = None if encoder_hidden_states is None else encoder_attention_mask + else: + # when attention_mask is defined: we don't even check for encoder_attention_mask. + # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. + # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. + # then we can simplify this whole if/else block to: + # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask + mask = attention_mask + + for resnet, attn in zip(self.resnets, self.attentions): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, temb, scale=lora_scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class KDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 4, + resnet_eps: float = 1e-5, + resnet_act_fn: str = "gelu", + resnet_group_size: int = 32, + add_downsample=False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + temb_channels=temb_channels, + groups=groups, + groups_out=groups_out, + eps=resnet_eps, + non_linearity=resnet_act_fn, + time_embedding_norm="ada_group", + conv_shortcut_bias=False, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + # YiYi's comments- might be able to use FirDownsample2D, look into details later + self.downsamplers = nn.ModuleList([KDownsample2D()]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None, scale: float = 1.0): + output_states = () + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale) + + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states, output_states + + +class KCrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + cross_attention_dim: int, + dropout: float = 0.0, + num_layers: int = 4, + resnet_group_size: int = 32, + add_downsample=True, + attention_head_dim: int = 64, + add_self_attention: bool = False, + resnet_eps: float = 1e-5, + resnet_act_fn: str = "gelu", + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + temb_channels=temb_channels, + groups=groups, + groups_out=groups_out, + eps=resnet_eps, + non_linearity=resnet_act_fn, + time_embedding_norm="ada_group", + conv_shortcut_bias=False, + ) + ) + attentions.append( + KAttentionBlock( + out_channels, + out_channels // attention_head_dim, + attention_head_dim, + cross_attention_dim=cross_attention_dim, + temb_channels=temb_channels, + attention_bias=True, + add_self_attention=add_self_attention, + cross_attention_norm="layer_norm", + group_size=resnet_group_size, + ) + ) + + self.resnets = nn.ModuleList(resnets) + self.attentions = nn.ModuleList(attentions) + + if add_downsample: + self.downsamplers = nn.ModuleList([KDownsample2D()]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + output_states = () + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + for resnet, attn in zip(self.resnets, self.attentions): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + + if self.downsamplers is None: + output_states += (None,) + else: + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states, output_states + + +class AttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + upsample_type="conv", + ): + super().__init__() + resnets = [] + attentions = [] + + self.upsample_type = upsample_type + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if upsample_type == "conv": + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + elif upsample_type == "resnet": + self.upsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + up=True, + ) + ] + ) + else: + self.upsamplers = None + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + for resnet, attn in zip(self.resnets, self.attentions): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb, scale=scale) + cross_attention_kwargs = {"scale": scale} + hidden_states = attn(hidden_states, **cross_attention_kwargs) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + if self.upsample_type == "resnet": + hidden_states = upsampler(hidden_states, temb=temb, scale=scale) + else: + hidden_states = upsampler(hidden_states, scale=scale) + + return hidden_states + + +class CrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + add_upsample=True, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + attention_type="default", + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + enter_block_number: Optional[int]=None, + ): + prv_f = [] + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)): + # pop res hidden states + + if enter_block_number is not None and i < len(self.resnets) - enter_block_number - 1: + continue + + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + prv_f.append(hidden_states) + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale) + + return hidden_states, prv_f + + +class UpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + i = 0 + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + i += 1 + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size, scale=scale) + + return hidden_states + + +class UpDecoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + temb_channels=None, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + def forward(self, hidden_states, temb=None, scale: float = 1.0): + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class AttnUpDecoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + add_upsample=True, + temb_channels=None, + ): + super().__init__() + resnets = [] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups if resnet_time_scale_shift != "spatial" else None, + spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + def forward(self, hidden_states, temb=None, scale: float = 1.0): + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb=temb, scale=scale) + cross_attention_kwargs = {"scale": scale} + hidden_states = attn(hidden_states, temb=temb, **cross_attention_kwargs) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, scale=scale) + + return hidden_states + + +class AttnSkipUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=np.sqrt(2.0), + add_upsample=True, + ): + super().__init__() + self.attentions = nn.ModuleList([]) + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + self.resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(resnet_in_channels + res_skip_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + self.attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=32, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) + if add_upsample: + self.resnet_up = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + up=True, + kernel="fir", + ) + self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + self.skip_norm = torch.nn.GroupNorm( + num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True + ) + self.act = nn.SiLU() + else: + self.resnet_up = None + self.skip_conv = None + self.skip_norm = None + self.act = None + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=None, scale: float = 1.0): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb, scale=scale) + + cross_attention_kwargs = {"scale": scale} + hidden_states = self.attentions[0](hidden_states, **cross_attention_kwargs) + + if skip_sample is not None: + skip_sample = self.upsampler(skip_sample) + else: + skip_sample = 0 + + if self.resnet_up is not None: + skip_sample_states = self.skip_norm(hidden_states) + skip_sample_states = self.act(skip_sample_states) + skip_sample_states = self.skip_conv(skip_sample_states) + + skip_sample = skip_sample + skip_sample_states + + hidden_states = self.resnet_up(hidden_states, temb, scale=scale) + + return hidden_states, skip_sample + + +class SkipUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_pre_norm: bool = True, + output_scale_factor=np.sqrt(2.0), + add_upsample=True, + upsample_padding=1, + ): + super().__init__() + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + self.resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min((resnet_in_channels + res_skip_channels) // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) + if add_upsample: + self.resnet_up = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + up=True, + kernel="fir", + ) + self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + self.skip_norm = torch.nn.GroupNorm( + num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True + ) + self.act = nn.SiLU() + else: + self.resnet_up = None + self.skip_conv = None + self.skip_norm = None + self.act = None + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=None, scale: float = 1.0): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb, scale=scale) + + if skip_sample is not None: + skip_sample = self.upsampler(skip_sample) + else: + skip_sample = 0 + + if self.resnet_up is not None: + skip_sample_states = self.skip_norm(hidden_states) + skip_sample_states = self.act(skip_sample_states) + skip_sample_states = self.skip_conv(skip_sample_states) + + skip_sample = skip_sample + skip_sample_states + + hidden_states = self.resnet_up(hidden_states, temb, scale=scale) + + return hidden_states, skip_sample + + +class ResnetUpsampleBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + skip_time_act=False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + up=True, + ) + ] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, temb, scale=scale) + + return hidden_states + + +class SimpleCrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + add_upsample=True, + skip_time_act=False, + only_cross_attention=False, + cross_attention_norm=None, + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.attention_head_dim = attention_head_dim + + self.num_heads = out_channels // self.attention_head_dim + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + processor = ( + AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() + ) + + attentions.append( + Attention( + query_dim=out_channels, + cross_attention_dim=out_channels, + heads=self.num_heads, + dim_head=self.attention_head_dim, + added_kv_proj_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + bias=True, + upcast_softmax=True, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + processor=processor, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + up=True, + ) + ] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + lora_scale = cross_attention_kwargs.get("scale", 1.0) + if attention_mask is None: + # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. + mask = None if encoder_hidden_states is None else encoder_attention_mask + else: + # when attention_mask is defined: we don't even check for encoder_attention_mask. + # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. + # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. + # then we can simplify this whole if/else block to: + # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask + mask = attention_mask + + for resnet, attn in zip(self.resnets, self.attentions): + # resnet + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, temb, scale=lora_scale) + + return hidden_states + + +class KUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 5, + resnet_eps: float = 1e-5, + resnet_act_fn: str = "gelu", + resnet_group_size: Optional[int] = 32, + add_upsample=True, + ): + super().__init__() + resnets = [] + k_in_channels = 2 * out_channels + k_out_channels = in_channels + num_layers = num_layers - 1 + + for i in range(num_layers): + in_channels = k_in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=k_out_channels if (i == num_layers - 1) else out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=groups, + groups_out=groups_out, + dropout=dropout, + non_linearity=resnet_act_fn, + time_embedding_norm="ada_group", + conv_shortcut_bias=False, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([KUpsample2D()]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + res_hidden_states_tuple = res_hidden_states_tuple[-1] + if res_hidden_states_tuple is not None: + hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1) + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class KCrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 4, + resnet_eps: float = 1e-5, + resnet_act_fn: str = "gelu", + resnet_group_size: int = 32, + attention_head_dim=1, # attention dim_head + cross_attention_dim: int = 768, + add_upsample: bool = True, + upcast_attention: bool = False, + ): + super().__init__() + resnets = [] + attentions = [] + + is_first_block = in_channels == out_channels == temb_channels + is_middle_block = in_channels != out_channels + add_self_attention = True if is_first_block else False + + self.has_cross_attention = True + self.attention_head_dim = attention_head_dim + + # in_channels, and out_channels for the block (k-unet) + k_in_channels = out_channels if is_first_block else 2 * out_channels + k_out_channels = in_channels + + num_layers = num_layers - 1 + + for i in range(num_layers): + in_channels = k_in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + if is_middle_block and (i == num_layers - 1): + conv_2d_out_channels = k_out_channels + else: + conv_2d_out_channels = None + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + conv_2d_out_channels=conv_2d_out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=groups, + groups_out=groups_out, + dropout=dropout, + non_linearity=resnet_act_fn, + time_embedding_norm="ada_group", + conv_shortcut_bias=False, + ) + ) + attentions.append( + KAttentionBlock( + k_out_channels if (i == num_layers - 1) else out_channels, + k_out_channels // attention_head_dim + if (i == num_layers - 1) + else out_channels // attention_head_dim, + attention_head_dim, + cross_attention_dim=cross_attention_dim, + temb_channels=temb_channels, + attention_bias=True, + add_self_attention=add_self_attention, + cross_attention_norm="layer_norm", + upcast_attention=upcast_attention, + ) + ) + + self.resnets = nn.ModuleList(resnets) + self.attentions = nn.ModuleList(attentions) + + if add_upsample: + self.upsamplers = nn.ModuleList([KUpsample2D()]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + res_hidden_states_tuple = res_hidden_states_tuple[-1] + if res_hidden_states_tuple is not None: + hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1) + + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + for resnet, attn in zip(self.resnets, self.attentions): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +# can potentially later be renamed to `No-feed-forward` attention +class KAttentionBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout: float = 0.0, + cross_attention_dim: Optional[int] = None, + attention_bias: bool = False, + upcast_attention: bool = False, + temb_channels: int = 768, # for ada_group_norm + add_self_attention: bool = False, + cross_attention_norm: Optional[str] = None, + group_size: int = 32, + ): + super().__init__() + self.add_self_attention = add_self_attention + + # 1. Self-Attn + if add_self_attention: + self.norm1 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size)) + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=None, + cross_attention_norm=None, + ) + + # 2. Cross-Attn + self.norm2 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size)) + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + cross_attention_norm=cross_attention_norm, + ) + + def _to_3d(self, hidden_states, height, weight): + return hidden_states.permute(0, 2, 3, 1).reshape(hidden_states.shape[0], height * weight, -1) + + def _to_4d(self, hidden_states, height, weight): + return hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, height, weight) + + def forward( + self, + hidden_states: torch.FloatTensor, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + # TODO: mark emb as non-optional (self.norm2 requires it). + # requires assessing impact of change to positional param interface. + emb: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + # 1. Self-Attention + if self.add_self_attention: + norm_hidden_states = self.norm1(hidden_states, emb) + + height, weight = norm_hidden_states.shape[2:] + norm_hidden_states = self._to_3d(norm_hidden_states, height, weight) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + attn_output = self._to_4d(attn_output, height, weight) + + hidden_states = attn_output + hidden_states + + # 2. Cross-Attention/None + norm_hidden_states = self.norm2(hidden_states, emb) + + height, weight = norm_hidden_states.shape[2:] + norm_hidden_states = self._to_3d(norm_hidden_states, height, weight) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask if encoder_hidden_states is None else encoder_attention_mask, + **cross_attention_kwargs, + ) + attn_output = self._to_4d(attn_output, height, weight) + + hidden_states = attn_output + hidden_states + + return hidden_states diff --git a/ixformer_sdk/contrib/DeepCache/sd/unet_2d_condition.py b/ixformer_sdk/contrib/DeepCache/sd/unet_2d_condition.py new file mode 100644 index 00000000..2b7ae7c0 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sd/unet_2d_condition.py @@ -0,0 +1,1257 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.utils.checkpoint + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import UNet2DConditionLoadersMixin +from diffusers.utils import BaseOutput, logging +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import ( + ADDED_KV_ATTENTION_PROCESSORS, + CROSS_ATTENTION_PROCESSORS, + AttentionProcessor, + AttnAddedKVProcessor, + AttnProcessor, +) +from diffusers.models.embeddings import ( + GaussianFourierProjection, + ImageHintTimeEmbedding, + ImageProjection, + ImageTimeEmbedding, + # PositionNet, + TextImageProjection, + TextImageTimeEmbedding, + TextTimeEmbedding, + TimestepEmbedding, + Timesteps, +) +from diffusers.models.modeling_utils import ModelMixin + +from .unet_2d_blocks import ( + UNetMidBlock2DCrossAttn, + UNetMidBlock2DSimpleCrossAttn, + get_down_block, + get_up_block, +) + +import time + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name +class FourierEmbedder(nn.Module): + def __init__(self, num_freqs=64, temperature=100): + super().__init__() + + self.num_freqs = num_freqs + self.temperature = temperature + + freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs) + freq_bands = freq_bands[None, None, None] + self.register_buffer("freq_bands", freq_bands, persistent=False) + + def __call__(self, x): + x = self.freq_bands * x.unsqueeze(-1) + return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 1, 3, 4, 2).reshape(*x.shape[:2], -1) + +class PositionNet(nn.Module): + def __init__(self, positive_len, out_dim, feature_type="text-only", fourier_freqs=8): + super().__init__() + self.positive_len = positive_len + self.out_dim = out_dim + + self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs) + self.position_dim = fourier_freqs * 2 * 4 # 2: sin/cos, 4: xyxy + + if isinstance(out_dim, tuple): + out_dim = out_dim[0] + + if feature_type == "text-only": + self.linears = nn.Sequential( + nn.Linear(self.positive_len + self.position_dim, 512), + nn.SiLU(), + nn.Linear(512, 512), + nn.SiLU(), + nn.Linear(512, out_dim), + ) + self.null_positive_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) + + elif feature_type == "text-image": + self.linears_text = nn.Sequential( + nn.Linear(self.positive_len + self.position_dim, 512), + nn.SiLU(), + nn.Linear(512, 512), + nn.SiLU(), + nn.Linear(512, out_dim), + ) + self.linears_image = nn.Sequential( + nn.Linear(self.positive_len + self.position_dim, 512), + nn.SiLU(), + nn.Linear(512, 512), + nn.SiLU(), + nn.Linear(512, out_dim), + ) + self.null_text_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) + self.null_image_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) + + self.null_position_feature = torch.nn.Parameter(torch.zeros([self.position_dim])) + + def forward( + self, + boxes, + masks, + positive_embeddings=None, + phrases_masks=None, + image_masks=None, + phrases_embeddings=None, + image_embeddings=None, + ): + masks = masks.unsqueeze(-1) + + # embedding position (it may includes padding as placeholder) + xyxy_embedding = self.fourier_embedder(boxes) # B*N*4 -> B*N*C + + # learnable null embedding + xyxy_null = self.null_position_feature.view(1, 1, -1) + + # replace padding with learnable null embedding + xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null + + # positionet with text only information + if positive_embeddings is not None: + # learnable null embedding + positive_null = self.null_positive_feature.view(1, 1, -1) + + # replace padding with learnable null embedding + positive_embeddings = positive_embeddings * masks + (1 - masks) * positive_null + + objs = self.linears(torch.cat([positive_embeddings, xyxy_embedding], dim=-1)) + + # positionet with text and image infomation + else: + phrases_masks = phrases_masks.unsqueeze(-1) + image_masks = image_masks.unsqueeze(-1) + + # learnable null embedding + text_null = self.null_text_feature.view(1, 1, -1) + image_null = self.null_image_feature.view(1, 1, -1) + + # replace padding with learnable null embedding + phrases_embeddings = phrases_embeddings * phrases_masks + (1 - phrases_masks) * text_null + image_embeddings = image_embeddings * image_masks + (1 - image_masks) * image_null + + objs_text = self.linears_text(torch.cat([phrases_embeddings, xyxy_embedding], dim=-1)) + objs_image = self.linears_image(torch.cat([image_embeddings, xyxy_embedding], dim=-1)) + objs = torch.cat([objs_text, objs_image], dim=1) + + return objs + +@dataclass +class UNet2DConditionOutput(BaseOutput): + """ + The output of [`UNet2DConditionModel`]. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. + """ + + sample: torch.FloatTensor = None + + +class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): + r""" + A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample + shaped output. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Parameters: + sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + Height and width of input/output sample. + in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample. + out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. + center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample. + flip_sin_to_cos (`bool`, *optional*, defaults to `False`): + Whether to flip the sin to cos in the time embedding. + freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. + down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): + The tuple of downsample blocks to use. + mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): + Block type for middle of UNet, it can be either `UNetMidBlock2DCrossAttn` or + `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): + The tuple of upsample blocks to use. + only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`): + Whether to include self-attention in the basic transformer blocks, see + [`~models.attention.BasicTransformerBlock`]. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + The tuple of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. + downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. + mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. + norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. + If `None`, normalization and activation layers is skipped in post-processing. + norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization. + cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + The dimension of the cross attention features. + transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for + [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], + [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. + encoder_hid_dim (`int`, *optional*, defaults to None): + If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` + dimension to `cross_attention_dim`. + encoder_hid_dim_type (`str`, *optional*, defaults to `None`): + If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text + embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`. + attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads. + num_attention_heads (`int`, *optional*): + The number of attention heads. If not defined, defaults to `attention_head_dim` + resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config + for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`. + class_embed_type (`str`, *optional*, defaults to `None`): + The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, + `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. + addition_embed_type (`str`, *optional*, defaults to `None`): + Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or + "text". "text" will use the `TextTimeEmbedding` layer. + addition_time_embed_dim: (`int`, *optional*, defaults to `None`): + Dimension for the timestep embeddings. + num_class_embeds (`int`, *optional*, defaults to `None`): + Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing + class conditioning with `class_embed_type` equal to `None`. + time_embedding_type (`str`, *optional*, defaults to `positional`): + The type of position embedding to use for timesteps. Choose from `positional` or `fourier`. + time_embedding_dim (`int`, *optional*, defaults to `None`): + An optional override for the dimension of the projected time embedding. + time_embedding_act_fn (`str`, *optional*, defaults to `None`): + Optional activation function to use only once on the time embeddings before they are passed to the rest of + the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`. + timestep_post_act (`str`, *optional*, defaults to `None`): + The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`. + time_cond_proj_dim (`int`, *optional*, defaults to `None`): + The dimension of `cond_proj` layer in the timestep embedding. + conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. + conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer. + projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when + `class_embed_type="projection"`. Required when `class_embed_type="projection"`. + class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time + embeddings with the class embeddings. + mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`): + Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If + `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the + `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False` + otherwise. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + center_input_sample: bool = False, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "DownBlock2D", + ), + mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", + up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: Union[int, Tuple[int]] = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + dropout: float = 0.0, + act_fn: str = "silu", + norm_num_groups: Optional[int] = 32, + norm_eps: float = 1e-5, + cross_attention_dim: Union[int, Tuple[int]] = 1280, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + encoder_hid_dim: Optional[int] = None, + encoder_hid_dim_type: Optional[str] = None, + attention_head_dim: Union[int, Tuple[int]] = 8, + num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + addition_embed_type: Optional[str] = None, + addition_time_embed_dim: Optional[int] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: int = 1.0, + time_embedding_type: str = "positional", + time_embedding_dim: Optional[int] = None, + time_embedding_act_fn: Optional[str] = None, + timestep_post_act: Optional[str] = None, + time_cond_proj_dim: Optional[int] = None, + conv_in_kernel: int = 3, + conv_out_kernel: int = 3, + projection_class_embeddings_input_dim: Optional[int] = None, + attention_type: str = "default", + class_embeddings_concat: bool = False, + mid_block_only_cross_attention: Optional[bool] = None, + cross_attention_norm: Optional[str] = None, + addition_embed_type_num_heads=64, + ): + super().__init__() + + self.sample_size = sample_size + + if num_attention_heads is not None: + raise ValueError( + "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." + ) + + # If `num_attention_heads` is not defined (which is the case for most models) + # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. + # The reason for this behavior is to correct for incorrectly named variables that were introduced + # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 + # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking + # which is why we correct for the naming here. + num_attention_heads = num_attention_heads or attention_head_dim + + # Check inputs + if len(down_block_types) != len(up_block_types): + raise ValueError( + f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." + ) + + if len(block_out_channels) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}." + ) + + if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." + ) + + # input + conv_in_padding = (conv_in_kernel - 1) // 2 + self.conv_in = nn.Conv2d( + in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding + ) + + # time + if time_embedding_type == "fourier": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 + if time_embed_dim % 2 != 0: + raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.") + self.time_proj = GaussianFourierProjection( + time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos + ) + timestep_input_dim = time_embed_dim + elif time_embedding_type == "positional": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 + + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) + timestep_input_dim = block_out_channels[0] + else: + raise ValueError( + f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." + ) + + self.time_embedding = TimestepEmbedding( + timestep_input_dim, + time_embed_dim, + act_fn=act_fn, + post_act_fn=timestep_post_act, + cond_proj_dim=time_cond_proj_dim, + ) + + if encoder_hid_dim_type is None and encoder_hid_dim is not None: + encoder_hid_dim_type = "text_proj" + self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type) + logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.") + + if encoder_hid_dim is None and encoder_hid_dim_type is not None: + raise ValueError( + f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}." + ) + + if encoder_hid_dim_type == "text_proj": + self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim) + elif encoder_hid_dim_type == "text_image_proj": + # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)` + self.encoder_hid_proj = TextImageProjection( + text_embed_dim=encoder_hid_dim, + image_embed_dim=cross_attention_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 + self.encoder_hid_proj = ImageProjection( + image_embed_dim=encoder_hid_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type is not None: + raise ValueError( + f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'." + ) + else: + self.encoder_hid_proj = None + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) + elif class_embed_type == "timestep": + self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn) + elif class_embed_type == "identity": + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + elif class_embed_type == "projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" + ) + # The projection `class_embed_type` is the same as the timestep `class_embed_type` except + # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings + # 2. it projects from an arbitrary input dimension. + # + # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. + # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. + # As a result, `TimestepEmbedding` can be passed arbitrary vectors. + self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) + elif class_embed_type == "simple_projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set" + ) + self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim) + else: + self.class_embedding = None + + if addition_embed_type == "text": + if encoder_hid_dim is not None: + text_time_embedding_from_dim = encoder_hid_dim + else: + text_time_embedding_from_dim = cross_attention_dim + + self.add_embedding = TextTimeEmbedding( + text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads + ) + elif addition_embed_type == "text_image": + # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)` + self.add_embedding = TextImageTimeEmbedding( + text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim + ) + elif addition_embed_type == "text_time": + self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) + self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) + elif addition_embed_type == "image": + # Kandinsky 2.2 + self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) + elif addition_embed_type == "image_hint": + # Kandinsky 2.2 ControlNet + self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) + elif addition_embed_type is not None: + raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.") + + if time_embedding_act_fn is None: + self.time_embed_act = None + else: + self.time_embed_act = get_activation(time_embedding_act_fn) + + self.down_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = only_cross_attention + + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = False + + if isinstance(num_attention_heads, int): + num_attention_heads = (num_attention_heads,) * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + if isinstance(cross_attention_dim, int): + cross_attention_dim = (cross_attention_dim,) * len(down_block_types) + + if isinstance(layers_per_block, int): + layers_per_block = [layers_per_block] * len(down_block_types) + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) + + if class_embeddings_concat: + # The time embeddings are concatenated with the class embeddings. The dimension of the + # time embeddings passed to the down, middle, and up blocks is twice the dimension of the + # regular time embeddings + blocks_time_embed_dim = time_embed_dim * 2 + else: + blocks_time_embed_dim = time_embed_dim + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block[i], + transformer_layers_per_block=transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + temb_channels=blocks_time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim[i], + num_attention_heads=num_attention_heads[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + dropout=dropout, + ) + self.down_blocks.append(down_block) + + # mid + if mid_block_type == "UNetMidBlock2DCrossAttn": + self.mid_block = UNetMidBlock2DCrossAttn( + transformer_layers_per_block=transformer_layers_per_block[-1], + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + dropout=dropout, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim[-1], + num_attention_heads=num_attention_heads[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn": + self.mid_block = UNetMidBlock2DSimpleCrossAttn( + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + dropout=dropout, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + cross_attention_dim=cross_attention_dim[-1], + attention_head_dim=attention_head_dim[-1], + resnet_groups=norm_num_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + only_cross_attention=mid_block_only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif mid_block_type is None: + self.mid_block = None + else: + raise ValueError(f"unknown mid_block_type : {mid_block_type}") + + # count how many layers upsample the images + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_num_attention_heads = list(reversed(num_attention_heads)) + reversed_layers_per_block = list(reversed(layers_per_block)) + reversed_cross_attention_dim = list(reversed(cross_attention_dim)) + reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) + only_cross_attention = list(reversed(only_cross_attention)) + + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=reversed_layers_per_block[i] + 1, + transformer_layers_per_block=reversed_transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=blocks_time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=reversed_cross_attention_dim[i], + num_attention_heads=reversed_num_attention_heads[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + dropout=dropout, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_num_groups is not None: + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps + ) + + self.conv_act = get_activation(act_fn) + + else: + self.conv_norm_out = None + self.conv_act = None + + conv_out_padding = (conv_out_kernel - 1) // 2 + self.conv_out = nn.Conv2d( + block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding + ) + + if attention_type in ["gated", "gated-text-image"]: + positive_len = 768 + if isinstance(cross_attention_dim, int): + positive_len = cross_attention_dim + elif isinstance(cross_attention_dim, tuple) or isinstance(cross_attention_dim, list): + positive_len = cross_attention_dim[0] + + feature_type = "text-only" if attention_type == "gated" else "text-image" + self.position_net = PositionNet( + positive_len=positive_len, out_dim=cross_attention_dim, feature_type=feature_type + ) + + @property + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attn_processor( + self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor, _remove_lora=_remove_lora) + else: + module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnAddedKVProcessor() + elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor, _remove_lora=True) + + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module splits the input tensor in slices to compute attention in + several steps. This is useful for saving some memory in exchange for a small decrease in speed. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If + `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): + if hasattr(module, "set_attention_slice"): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_sliceable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_sliceable_dims(module) + + num_sliceable_layers = len(sliceable_head_dims) + + if slice_size == "auto": + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == "max": + # make smallest slice possible + slice_size = num_sliceable_layers * [1] + + slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" + f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError(f"size {size} has to be smaller or equal to {dim}.") + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + if hasattr(module, "set_attention_slice"): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + class_labels: Optional[torch.Tensor] = None, + timestep_cond: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + mid_block_additional_residual: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + quick_replicate: bool = False, + replicate_prv_feature: Optional[List[torch.Tensor]] = None, + cache_layer_id: Optional[int] = None, + cache_block_id: Optional[int] = None, + return_dict: bool = True, + ) -> Union[UNet2DConditionOutput, Tuple]: + r""" + The [`UNet2DConditionModel`] forward method. + + Args: + sample (`torch.FloatTensor`): + The noisy input tensor with the following shape `(batch, channel, height, width)`. + timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. + encoder_hidden_states (`torch.FloatTensor`): + The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. + encoder_attention_mask (`torch.Tensor`): + A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If + `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, + which adds large negative values to the attention scores corresponding to "discard" tokens. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + + Returns: + [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: + If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise + a `tuple` is returned where the first element is the sample tensor. + """ + # By default samples have to be AT least a multiple of the overall upsampling factor. + # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). + # However, the upsampling interpolation output size can be forced to fit any upsampling size + # on the fly if necessary. + default_overall_up_factor = 2**self.num_upsamplers + + # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` + forward_upsample_size = False + upsample_size = None + + if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): + logger.info("Forward upsample size to force interpolation output size.") + forward_upsample_size = True + + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension + # expects mask of shape: + # [batch, key_tokens] + # adds singleton query_tokens dimension: + # [batch, 1, key_tokens] + # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: + # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) + # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) + if attention_mask is not None: + # assume that mask is expressed as: + # (1 = keep, 0 = discard) + # convert mask into a bias that can be added to attention scores: + # (keep = +0, discard = -10000.0) + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # convert encoder_attention_mask to a bias the same way we do for attention_mask + if encoder_attention_mask is not None: + encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # 0. center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb, timestep_cond) + aug_emb = None + + if self.class_embedding is not None: + if class_labels is None: + raise ValueError("class_labels should be provided when num_class_embeds > 0") + + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # there might be better ways to encapsulate this. + class_labels = class_labels.to(dtype=sample.dtype) + + class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype) + + if self.config.class_embeddings_concat: + emb = torch.cat([emb, class_emb], dim=-1) + else: + emb = emb + class_emb + + if self.config.addition_embed_type == "text": + aug_emb = self.add_embedding(encoder_hidden_states) + elif self.config.addition_embed_type == "text_image": + # Kandinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + + image_embs = added_cond_kwargs.get("image_embeds") + text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states) + aug_emb = self.add_embedding(text_embs, image_embs) + elif self.config.addition_embed_type == "text_time": + # SDXL - style + if "text_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" + ) + text_embeds = added_cond_kwargs.get("text_embeds") + if "time_ids" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" + ) + time_ids = added_cond_kwargs.get("time_ids") + time_embeds = self.add_time_proj(time_ids.flatten()) + time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) + + add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) + add_embeds = add_embeds.to(emb.dtype) + aug_emb = self.add_embedding(add_embeds) + elif self.config.addition_embed_type == "image": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + aug_emb = self.add_embedding(image_embs) + elif self.config.addition_embed_type == "image_hint": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + hint = added_cond_kwargs.get("hint") + aug_emb, hint = self.add_embedding(image_embs, hint) + sample = torch.cat([sample, hint], dim=1) + + emb = emb + aug_emb if aug_emb is not None else emb + + if self.time_embed_act is not None: + emb = self.time_embed_act(emb) + + if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj": + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj": + # Kadinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(image_embeds) + # 2. pre-process + sample = self.conv_in(sample) + + # 2.5 GLIGEN position net + if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None: + cross_attention_kwargs = cross_attention_kwargs.copy() + gligen_args = cross_attention_kwargs.pop("gligen") + cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)} + + # 3. down + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None + is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None + + down_block_res_samples = (sample,) + if quick_replicate and replicate_prv_feature is not None: + # Down + for i, downsample_block in enumerate(self.down_blocks): + if i > cache_layer_id: + break + + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + # For t2i-adapter CrossAttnDownBlock2D + additional_residuals = {} + if is_adapter and len(down_block_additional_residuals) > 0: + additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0) + + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + exist_block_number=cache_block_id if i == cache_layer_id else None, + **additional_residuals, + ) + else: + sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale) + + if is_adapter and len(down_block_additional_residuals) > 0: + sample += down_block_additional_residuals.pop(0) + + down_block_res_samples += res_samples + + # No Middle + # Up + #print("down_block_res_samples:", [res_sample.shape for res_sample in down_block_res_samples]) + sample = replicate_prv_feature + #down_block_res_samples = down_block_res_samples[:-1] + if cache_block_id == len(self.down_blocks[cache_layer_id].attentions) : + cache_block_id = 0 + cache_layer_id += 1 + else: + cache_block_id += 1 + + for i, upsample_block in enumerate(self.up_blocks): + if i < len(self.up_blocks) - 1 - cache_layer_id: + continue + + if i == len(self.up_blocks) - 1 - cache_layer_id: + trunc_upsample_block = cache_block_id + 1 + else: + trunc_upsample_block = len(upsample_block.resnets) + + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-trunc_upsample_block:] + down_block_res_samples = down_block_res_samples[: -trunc_upsample_block] + + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + #print(sample.shape, [res_sample.shape for res_sample in res_samples]) + sample, _ = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + enter_block_number=cache_block_id if i == len(self.up_blocks) - 1 - cache_layer_id else None, + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size, + scale=lora_scale, + ) + + prv_f = replicate_prv_feature + else: + for i, downsample_block in enumerate(self.down_blocks): + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + # For t2i-adapter CrossAttnDownBlock2D + additional_residuals = {} + if is_adapter and len(down_block_additional_residuals) > 0: + additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0) + + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + **additional_residuals, + ) + else: + sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale) + + if is_adapter and len(down_block_additional_residuals) > 0: + sample += down_block_additional_residuals.pop(0) + + down_block_res_samples += res_samples + + if is_controlnet: + new_down_block_res_samples = () + + for down_block_res_sample, down_block_additional_residual in zip( + down_block_res_samples, down_block_additional_residuals + ): + down_block_res_sample = down_block_res_sample + down_block_additional_residual + new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) + + down_block_res_samples = new_down_block_res_samples + + # 4. mid + if self.mid_block is not None: + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + # To support T2I-Adapter-XL + if ( + is_adapter + and len(down_block_additional_residuals) > 0 + and sample.shape == down_block_additional_residuals[0].shape + ): + sample += down_block_additional_residuals.pop(0) + + if is_controlnet: + sample = sample + mid_block_additional_residual + + # 5. up + if cache_block_id is not None: + if cache_block_id == len(self.down_blocks[cache_layer_id].attentions) : + cache_block_id = 0 + cache_layer_id += 1 + else: + cache_block_id += 1 + #print("down_block_res_samples:", [res_sample.shape for res_sample in down_block_res_samples]) + #print(cache_block_id, cache_layer_id) + prv_f = None + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] + #print(sample.shape, [res_sample.shape for res_sample in res_samples]) + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + sample, current_record_f = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size, + scale=lora_scale, + ) + current_record_f = None + + #print("Append prv_feature with shape:", sample.shape) + if cache_layer_id is not None and current_record_f is not None and i == len(self.up_blocks) - cache_layer_id - 1: + prv_f = current_record_f[-cache_block_id-1] + + # 6. post-process + if self.conv_norm_out: + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + if not return_dict: + return (sample, prv_f,) + + return UNet2DConditionOutput(sample=sample) diff --git a/ixformer_sdk/contrib/DeepCache/sdxl/__init__.py b/ixformer_sdk/contrib/DeepCache/sdxl/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/DeepCache/sdxl/pipeline_stable_diffusion_xl.py b/ixformer_sdk/contrib/DeepCache/sdxl/pipeline_stable_diffusion_xl.py new file mode 100644 index 00000000..d0e9f4ac --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sdxl/pipeline_stable_diffusion_xl.py @@ -0,0 +1,1100 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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 inspect +import os +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +import torch +from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer + +from diffusers.image_processor import VaeImageProcessor +from diffusers.loaders import ( + FromSingleFileMixin, + LoraLoaderMixin, + TextualInversionLoaderMixin, +) +from diffusers.models import AutoencoderKL +from diffusers.models.attention_processor import ( + AttnProcessor2_0, + LoRAAttnProcessor2_0, + LoRAXFormersAttnProcessor, + XFormersAttnProcessor, +) +from diffusers.models.lora import adjust_lora_scale_text_encoder +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils import ( + is_accelerate_available, + is_accelerate_version, + is_invisible_watermark_available, + logging, + replace_example_docstring, +) +from diffusers.utils.torch_utils import randn_tensor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput + +from .unet_2d_condition import UNet2DConditionModel +from .pipeline_utils import DiffusionPipeline + +if is_invisible_watermark_available(): + from diffusers.pipelines.stable_diffusion_xl.watermark import StableDiffusionXLWatermarker + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import StableDiffusionXLPipeline + + >>> pipe = StableDiffusionXLPipeline.from_pretrained( + ... "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 + ... ) + >>> pipe = pipe.to("cuda") + + >>> prompt = "a photo of an astronaut riding a horse on mars" + >>> image = pipe(prompt).images[0] + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + """ + Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + return noise_cfg + +def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2): + if pow is None: + pow=1.2 + if center is None: + center=0 + import numpy as np + while pow > 1: + # Generate linearly spaced values between 0 and a max value + x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1) + indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]] + if len(indices) == n_samples: + break + pow -=0.02 + if pow <= 1: + raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.") + return indices, pow + +class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin): + r""" + Pipeline for text-to-image generation using Stable Diffusion XL. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the + library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) + + In addition the pipeline inherits the following loading methods: + - *LoRA*: [`StableDiffusionXLPipeline.load_lora_weights`] + - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`] + + as well as the following saving methods: + - *LoRA*: [`loaders.StableDiffusionXLPipeline.save_lora_weights`] + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`CLIPTextModel`]): + Frozen text-encoder. Stable Diffusion XL uses the text portion of + [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically + the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. + text_encoder_2 ([` CLIPTextModelWithProjection`]): + Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of + [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), + specifically the + [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k) + variant. + tokenizer (`CLIPTokenizer`): + Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). + tokenizer_2 (`CLIPTokenizer`): + Second Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). + unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of + [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. + force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`): + Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of + `stabilityai/stable-diffusion-xl-base-1-0`. + add_watermarker (`bool`, *optional*): + Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to + watermark output images. If not defined, it will default to True if the package is installed, otherwise no + watermarker will be used. + """ + model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + text_encoder_2: CLIPTextModelWithProjection, + tokenizer: CLIPTokenizer, + tokenizer_2: CLIPTokenizer, + unet: UNet2DConditionModel, + scheduler: KarrasDiffusionSchedulers, + force_zeros_for_empty_prompt: bool = True, + add_watermarker: Optional[bool] = None, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + unet=unet, + scheduler=scheduler, + ) + self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.default_sample_size = self.unet.config.sample_size + + add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() + + if add_watermarker: + self.watermark = StableDiffusionXLWatermarker() + else: + self.watermark = None + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.vae.enable_slicing() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_slicing() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_tiling + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + self.vae.enable_tiling() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_tiling + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_tiling() + + def encode_prompt( + self, + prompt: str, + prompt_2: Optional[str] = None, + device: Optional[torch.device] = None, + num_images_per_prompt: int = 1, + do_classifier_free_guidance: bool = True, + negative_prompt: Optional[str] = None, + negative_prompt_2: Optional[str] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + lora_scale: Optional[float] = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + used in both text-encoders + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + do_classifier_free_guidance (`bool`): + whether to use classifier free guidance or not + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + negative_prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and + `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `prompt` input argument. + negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` + input argument. + lora_scale (`float`, *optional*): + A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. + """ + device = device or self._execution_device + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, LoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + adjust_lora_scale_text_encoder(self.text_encoder, lora_scale) + adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale) + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # Define tokenizers and text encoders + tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2] + text_encoders = ( + [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2] + ) + + if prompt_embeds is None: + prompt_2 = prompt_2 or prompt + # textual inversion: procecss multi-vector tokens if necessary + prompt_embeds_list = [] + prompts = [prompt, prompt_2] + for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, tokenizer) + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids + untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( + text_input_ids, untruncated_ids + ): + removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {tokenizer.model_max_length} tokens: {removed_text}" + ) + + prompt_embeds = text_encoder( + text_input_ids.to(device), + output_hidden_states=True, + ) + + # We are only ALWAYS interested in the pooled output of the final text encoder + pooled_prompt_embeds = prompt_embeds[0] + prompt_embeds = prompt_embeds.hidden_states[-2] + + prompt_embeds_list.append(prompt_embeds) + + prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) + + # get unconditional embeddings for classifier free guidance + zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt + if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt: + negative_prompt_embeds = torch.zeros_like(prompt_embeds) + negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds) + elif do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt_2 = negative_prompt_2 or negative_prompt + + uncond_tokens: List[str] + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt, negative_prompt_2] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = [negative_prompt, negative_prompt_2] + + negative_prompt_embeds_list = [] + for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders): + if isinstance(self, TextualInversionLoaderMixin): + negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer) + + max_length = prompt_embeds.shape[1] + uncond_input = tokenizer( + negative_prompt, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + negative_prompt_embeds = text_encoder( + uncond_input.input_ids.to(device), + output_hidden_states=True, + ) + # We are only ALWAYS interested in the pooled output of the final text encoder + negative_pooled_prompt_embeds = negative_prompt_embeds[0] + negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2] + + negative_prompt_embeds_list.append(negative_prompt_embeds) + + negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1) + + prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( + bs_embed * num_images_per_prompt, -1 + ) + if do_classifier_free_guidance: + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( + bs_embed * num_images_per_prompt, -1 + ) + + return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs( + self, + prompt, + prompt_2, + height, + width, + callback_steps, + negative_prompt=None, + negative_prompt_2=None, + prompt_embeds=None, + negative_prompt_embeds=None, + pooled_prompt_embeds=None, + negative_pooled_prompt_embeds=None, + ): + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt_2 is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): + raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_2 is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + + if prompt_embeds is not None and pooled_prompt_embeds is None: + raise ValueError( + "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`." + ) + + if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None: + raise ValueError( + "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`." + ) + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): + shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + return latents + + def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype): + add_time_ids = list(original_size + crops_coords_top_left + target_size) + + passed_add_embed_dim = ( + self.unet.config.addition_time_embed_dim * len(add_time_ids) + self.text_encoder_2.config.projection_dim + ) + expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features + + if expected_add_embed_dim != passed_add_embed_dim: + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." + ) + + add_time_ids = torch.tensor([add_time_ids], dtype=dtype) + return add_time_ids + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae + def upcast_vae(self): + dtype = self.vae.dtype + self.vae.to(dtype=torch.float32) + use_torch_2_0_or_xformers = isinstance( + self.vae.decoder.mid_block.attentions[0].processor, + ( + AttnProcessor2_0, + XFormersAttnProcessor, + LoRAXFormersAttnProcessor, + LoRAAttnProcessor2_0, + ), + ) + # if xformers or torch_2_0 is used attention block does not need + # to be in float32 which can save lots of memory + if use_torch_2_0_or_xformers: + self.vae.post_quant_conv.to(dtype) + self.vae.decoder.conv_in.to(dtype) + self.vae.decoder.mid_block.to(dtype) + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + prompt_2: Optional[Union[str, List[str]]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + denoising_end: Optional[float] = None, + guidance_scale: float = 5.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt_2: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: int = 1, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guidance_rescale: float = 0.0, + original_size: Optional[Tuple[int, int]] = None, + crops_coords_top_left: Tuple[int, int] = (0, 0), + target_size: Optional[Tuple[int, int]] = None, + negative_original_size: Optional[Tuple[int, int]] = None, + negative_crops_coords_top_left: Tuple[int, int] = (0, 0), + negative_target_size: Optional[Tuple[int, int]] = None, + cache_interval: int = 1, + cache_layer_id: int = None, + cache_block_id: int = None, + uniform: bool = True, + pow: float = None, + center: int = None, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + used in both text-encoders + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + Anything below 512 pixels won't work well for + [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) + and checkpoints that are not specifically fine-tuned on low resolutions. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + Anything below 512 pixels won't work well for + [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) + and checkpoints that are not specifically fine-tuned on low resolutions. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + denoising_end (`float`, *optional*): + When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be + completed before it is intentionally prematurely terminated. As a result, the returned sample will + still retain a substantial amount of noise as determined by the discrete timesteps selected by the + scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a + "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image + Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) + guidance_scale (`float`, *optional*, defaults to 5.0): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + negative_prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and + `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to + [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `prompt` input argument. + negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` + input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead + of a plain tuple. + callback (`Callable`, *optional*): + A function that will be called every `callback_steps` steps during inference. The function will be + called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function will be called. If not specified, the callback will be + called at every step. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + guidance_rescale (`float`, *optional*, defaults to 0.7): + Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of + [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). + Guidance rescale factor should fix overexposure when using zero terminal SNR. + original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. + `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as + explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position + `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting + `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + For most cases, `target_size` should be set to the desired height and width of the generated image. If + not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in + section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + To negatively condition the generation process based on a specific image resolution. Part of SDXL's + micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more + information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. + negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's + micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more + information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. + negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + To negatively condition the generation process based on a target image resolution. It should be as same + as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more + information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. + + Examples: + + Returns: + [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`: + [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a + `tuple`. When returning a tuple, the first element is a list with the generated images. + """ + # 0. Default height and width to unet + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + original_size = original_size or (height, width) + target_size = target_size or (height, width) + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + prompt_2, + height, + width, + callback_steps, + negative_prompt, + negative_prompt_2, + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + text_encoder_lora_scale = ( + cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None + ) + ( + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + device=device, + num_images_per_prompt=num_images_per_prompt, + do_classifier_free_guidance=do_classifier_free_guidance, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + lora_scale=text_encoder_lora_scale, + ) + + # 4. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + + timesteps = self.scheduler.timesteps + + # 5. Prepare latent variables + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7. Prepare added time ids & embeddings + add_text_embeds = pooled_prompt_embeds + add_time_ids = self._get_add_time_ids( + original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype + ) + if negative_original_size is not None and negative_target_size is not None: + negative_add_time_ids = self._get_add_time_ids( + negative_original_size, + negative_crops_coords_top_left, + negative_target_size, + dtype=prompt_embeds.dtype, + ) + else: + negative_add_time_ids = add_time_ids + + if do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) + add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0) + + prompt_embeds = prompt_embeds.to(device) + add_text_embeds = add_text_embeds.to(device) + add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1) + + # 8. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + + # 7.1 Apply denoising_end + if denoising_end is not None and isinstance(denoising_end, float) and denoising_end > 0 and denoising_end < 1: + discrete_timestep_cutoff = int( + round( + self.scheduler.config.num_train_timesteps + - (denoising_end * self.scheduler.config.num_train_timesteps) + ) + ) + num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) + timesteps = timesteps[:num_inference_steps] + + if cache_interval == 1: + interval_seq = list(range(num_inference_steps)) + else: + if uniform: + interval_seq = list(range(0, num_inference_steps, cache_interval)) + else: + num_slow_step = num_inference_steps//cache_interval + if num_inference_steps%cache_interval != 0: + num_slow_step += 1 + + interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,] + #print(interval_seq) + + prv_features = None + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + if i in interval_seq: + prv_features = None + #print(t, prv_features is None) + + # predict the noise residual + added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} + # print(f"{latent_model_input.shape},{t.shape},{prompt_embeds.shape} {cross_attention_kwargs} {added_cond_kwargs.keys()} {added_cond_kwargs['text_embeds'].shape} {added_cond_kwargs['time_ids'].shape} {prv_features.shape if prv_features is not None else None }" ) + + if not latent_model_input.requires_grad and latent_model_input.shape[0]>=128:#add for xl batch64 oom,batch*2为unet 输入的batch + _chunk_size=int(os.environ.get("ENABLE_IXFORMER_UNET_CHUNKSIZE", "8"))#unet 输入batch chunk为8时不oom + num_chunks = latent_model_input.shape[0] // _chunk_size + noise_pred_list=[] + prv_features_list=[] + for latent_model_input_slice,\ + prompt_embeds_slice,added_cond_kwargs_slice_text_embeds,\ + added_cond_kwargs_slice_time_ids,prv_features_slice\ + in zip(latent_model_input.chunk(num_chunks, dim=0), + prompt_embeds.chunk(num_chunks, dim=0), + added_cond_kwargs["text_embeds"].chunk(num_chunks, dim=0), + added_cond_kwargs["time_ids"].chunk(num_chunks, dim=0), + prv_features.chunk(num_chunks, dim=0) if prv_features is not None else [None] *num_chunks + ): + added_cond_kwargs_slice={} + added_cond_kwargs_slice["text_embeds"]=added_cond_kwargs_slice_text_embeds + added_cond_kwargs_slice["time_ids"]=added_cond_kwargs_slice_time_ids + noise_pred_item, prv_features_item = self.unet( + latent_model_input_slice, + t, + encoder_hidden_states=prompt_embeds_slice, + cross_attention_kwargs=cross_attention_kwargs, + added_cond_kwargs=added_cond_kwargs_slice, + replicate_prv_feature=prv_features_slice, + quick_replicate= cache_interval>1, + cache_layer_id=cache_layer_id, + cache_block_id=cache_block_id, + return_dict=False, + ) + noise_pred_list.append(noise_pred_item) + prv_features_list.append(prv_features_item) + noise_pred = torch.cat(noise_pred_list,dim=0) + prv_features = torch.cat(prv_features_list,dim=0) + del noise_pred_list + del prv_features_list + torch.cuda.empty_cache() + + else: + noise_pred, prv_features = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + added_cond_kwargs=added_cond_kwargs, + replicate_prv_feature=prv_features, + quick_replicate= cache_interval>1, + cache_layer_id=cache_layer_id, + cache_block_id=cache_block_id, + return_dict=False, + ) + + + + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + if do_classifier_free_guidance and guidance_rescale > 0.0: + # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf + noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + callback(i, t, latents) + + if not output_type == "latent": + # make sure the VAE is in float32 mode, as it overflows in float16 + needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast + + if needs_upcasting: + self.upcast_vae() + latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) + batch_number = latents.shape[0] + rebatching=False + + if latents.shape[-1] >=128:#1024 + chunck_size =int(os.environ.get("ENABLE_IXFORMER_VAE_CHUNKSIZE", "4")) + batch_number = chunck_size if batch_number>chunck_size else batch_number#1024x1024 batch >4 oom + rebatching=True + elif latents.shape[-1] >=64:#512 + chunck_size =int(os.environ.get("ENABLE_IXFORMER_VAE_CHUNKSIZE", "8")) + batch_number = chunck_size if batch_number>chunck_size else batch_number#512 batch >8 oom + rebatching=True + if not rebatching: + image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] + else: + image = torch.empty((latents.shape[0],3,height, width), device=latents.device) + for x in range(0, latents.shape[0], batch_number): + batch_end =min(x+batch_number,latents.shape[0]) + latents_each = latents[x:batch_end] + + + image[x:batch_end] = self.vae.decode(latents_each / self.vae.config.scaling_factor, return_dict=False)[0] + del latents_each + torch.cuda.empty_cache() + + + # cast back to fp16 if needed + if needs_upcasting: + self.vae.to(dtype=torch.float16) + else: + image = latents + + if not output_type == "latent": + # apply watermark if available + if self.watermark is not None: + image = self.watermark.apply_watermark(image) + + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return StableDiffusionXLPipelineOutput(images=image) + + # Overrride to properly handle the loading and unloading of the additional text encoder. + def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs): + # We could have accessed the unet config from `lora_state_dict()` too. We pass + # it here explicitly to be able to tell that it's coming from an SDXL + # pipeline. + + # Remove any existing hooks. + if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"): + from accelerate.hooks import AlignDevicesHook, CpuOffload, remove_hook_from_module + else: + raise ImportError("Offloading requires `accelerate v0.17.0` or higher.") + + is_model_cpu_offload = False + is_sequential_cpu_offload = False + recursive = False + for _, component in self.components.items(): + if isinstance(component, torch.nn.Module): + if hasattr(component, "_hf_hook"): + is_model_cpu_offload = isinstance(getattr(component, "_hf_hook"), CpuOffload) + is_sequential_cpu_offload = isinstance(getattr(component, "_hf_hook"), AlignDevicesHook) + logger.info( + "Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again." + ) + recursive = is_sequential_cpu_offload + remove_hook_from_module(component, recurse=recursive) + state_dict, network_alphas = self.lora_state_dict( + pretrained_model_name_or_path_or_dict, + unet_config=self.unet.config, + **kwargs, + ) + self.load_lora_into_unet(state_dict, network_alphas=network_alphas, unet=self.unet) + + text_encoder_state_dict = {k: v for k, v in state_dict.items() if "text_encoder." in k} + if len(text_encoder_state_dict) > 0: + self.load_lora_into_text_encoder( + text_encoder_state_dict, + network_alphas=network_alphas, + text_encoder=self.text_encoder, + prefix="text_encoder", + lora_scale=self.lora_scale, + ) + + text_encoder_2_state_dict = {k: v for k, v in state_dict.items() if "text_encoder_2." in k} + if len(text_encoder_2_state_dict) > 0: + self.load_lora_into_text_encoder( + text_encoder_2_state_dict, + network_alphas=network_alphas, + text_encoder=self.text_encoder_2, + prefix="text_encoder_2", + lora_scale=self.lora_scale, + ) + + # Offload back. + if is_model_cpu_offload: + self.enable_model_cpu_offload() + elif is_sequential_cpu_offload: + self.enable_sequential_cpu_offload() + + @classmethod + def save_lora_weights( + self, + save_directory: Union[str, os.PathLike], + unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + is_main_process: bool = True, + weight_name: str = None, + save_function: Callable = None, + safe_serialization: bool = True, + ): + state_dict = {} + + def pack_weights(layers, prefix): + layers_weights = layers.state_dict() if isinstance(layers, torch.nn.Module) else layers + layers_state_dict = {f"{prefix}.{module_name}": param for module_name, param in layers_weights.items()} + return layers_state_dict + + if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers): + raise ValueError( + "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers` or `text_encoder_2_lora_layers`." + ) + + if unet_lora_layers: + state_dict.update(pack_weights(unet_lora_layers, "unet")) + + if text_encoder_lora_layers and text_encoder_2_lora_layers: + state_dict.update(pack_weights(text_encoder_lora_layers, "text_encoder")) + state_dict.update(pack_weights(text_encoder_2_lora_layers, "text_encoder_2")) + + self.write_lora_layers( + state_dict=state_dict, + save_directory=save_directory, + is_main_process=is_main_process, + weight_name=weight_name, + save_function=save_function, + safe_serialization=safe_serialization, + ) + + def _remove_text_encoder_monkey_patch(self): + self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder) + self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder_2) diff --git a/ixformer_sdk/contrib/DeepCache/sdxl/pipeline_stable_diffusion_xl_img2img.py b/ixformer_sdk/contrib/DeepCache/sdxl/pipeline_stable_diffusion_xl_img2img.py new file mode 100644 index 00000000..1746eaa2 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sdxl/pipeline_stable_diffusion_xl_img2img.py @@ -0,0 +1,1187 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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 inspect +import os +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import PIL.Image +import torch +from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer + +from diffusers.image_processor import PipelineImageInput, VaeImageProcessor +from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin +from diffusers.models import AutoencoderKL, UNet2DConditionModel +from diffusers.models.attention_processor import ( + AttnProcessor2_0, + LoRAAttnProcessor2_0, + LoRAXFormersAttnProcessor, + XFormersAttnProcessor, +) +from diffusers.models.lora import adjust_lora_scale_text_encoder +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils import ( + is_accelerate_available, + is_accelerate_version, + is_invisible_watermark_available, + logging, + replace_example_docstring, +) +from diffusers.utils.torch_utils import randn_tensor +from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput + +from .pipeline_utils import DiffusionPipeline + + +if is_invisible_watermark_available(): + from diffusers.pipelines.stable_diffusion_xl.watermark import StableDiffusionXLWatermarker + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import StableDiffusionXLImg2ImgPipeline + >>> from diffusers.utils import load_image + + >>> pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained( + ... "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16 + ... ) + >>> pipe = pipe.to("cuda") + >>> url = "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png" + + >>> init_image = load_image(url).convert("RGB") + >>> prompt = "a photo of an astronaut riding a horse on mars" + >>> image = pipe(prompt, image=init_image).images[0] + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + """ + Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + return noise_cfg + +def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2): + if pow is None: + pow=1.2 + if center is None: + center=0 + import numpy as np + while pow > 1: + # Generate linearly spaced values between 0 and a max value + x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1) + indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]] + if len(indices) == n_samples: + break + pow -=0.02 + if pow <= 1: + raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.") + return indices, pow + +class StableDiffusionXLImg2ImgPipeline( + DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin +): + r""" + Pipeline for text-to-image generation using Stable Diffusion XL. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the + library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) + + In addition the pipeline inherits the following loading methods: + - *LoRA*: [`loaders.LoraLoaderMixin.load_lora_weights`] + - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`] + + as well as the following saving methods: + - *LoRA*: [`loaders.LoraLoaderMixin.save_lora_weights`] + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`CLIPTextModel`]): + Frozen text-encoder. Stable Diffusion XL uses the text portion of + [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically + the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. + text_encoder_2 ([` CLIPTextModelWithProjection`]): + Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of + [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), + specifically the + [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k) + variant. + tokenizer (`CLIPTokenizer`): + Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). + tokenizer_2 (`CLIPTokenizer`): + Second Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). + unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of + [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. + requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`): + Whether the `unet` requires an `aesthetic_score` condition to be passed during inference. Also see the + config of `stabilityai/stable-diffusion-xl-refiner-1-0`. + force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`): + Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of + `stabilityai/stable-diffusion-xl-base-1-0`. + add_watermarker (`bool`, *optional*): + Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to + watermark output images. If not defined, it will default to True if the package is installed, otherwise no + watermarker will be used. + """ + model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" + + _optional_components = ["tokenizer", "text_encoder"] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + text_encoder_2: CLIPTextModelWithProjection, + tokenizer: CLIPTokenizer, + tokenizer_2: CLIPTokenizer, + unet: UNet2DConditionModel, + scheduler: KarrasDiffusionSchedulers, + requires_aesthetics_score: bool = False, + force_zeros_for_empty_prompt: bool = True, + add_watermarker: Optional[bool] = None, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + unet=unet, + scheduler=scheduler, + ) + self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) + self.register_to_config(requires_aesthetics_score=requires_aesthetics_score) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + + add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available() + + if add_watermarker: + self.watermark = StableDiffusionXLWatermarker() + else: + self.watermark = None + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.vae.enable_slicing() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_slicing() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_tiling + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + self.vae.enable_tiling() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_tiling + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_tiling() + + # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt + def encode_prompt( + self, + prompt: str, + prompt_2: Optional[str] = None, + device: Optional[torch.device] = None, + num_images_per_prompt: int = 1, + do_classifier_free_guidance: bool = True, + negative_prompt: Optional[str] = None, + negative_prompt_2: Optional[str] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + lora_scale: Optional[float] = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + used in both text-encoders + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + do_classifier_free_guidance (`bool`): + whether to use classifier free guidance or not + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + negative_prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and + `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `prompt` input argument. + negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` + input argument. + lora_scale (`float`, *optional*): + A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. + """ + device = device or self._execution_device + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, LoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + adjust_lora_scale_text_encoder(self.text_encoder, lora_scale) + adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale) + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # Define tokenizers and text encoders + tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2] + text_encoders = ( + [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2] + ) + + if prompt_embeds is None: + prompt_2 = prompt_2 or prompt + # textual inversion: procecss multi-vector tokens if necessary + prompt_embeds_list = [] + prompts = [prompt, prompt_2] + for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders): + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, tokenizer) + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids + untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( + text_input_ids, untruncated_ids + ): + removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {tokenizer.model_max_length} tokens: {removed_text}" + ) + + prompt_embeds = text_encoder( + text_input_ids.to(device), + output_hidden_states=True, + ) + + # We are only ALWAYS interested in the pooled output of the final text encoder + pooled_prompt_embeds = prompt_embeds[0] + prompt_embeds = prompt_embeds.hidden_states[-2] + + prompt_embeds_list.append(prompt_embeds) + + prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) + + # get unconditional embeddings for classifier free guidance + zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt + if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt: + negative_prompt_embeds = torch.zeros_like(prompt_embeds) + negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds) + elif do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt_2 = negative_prompt_2 or negative_prompt + + uncond_tokens: List[str] + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt, negative_prompt_2] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = [negative_prompt, negative_prompt_2] + + negative_prompt_embeds_list = [] + for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders): + if isinstance(self, TextualInversionLoaderMixin): + negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer) + + max_length = prompt_embeds.shape[1] + uncond_input = tokenizer( + negative_prompt, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + negative_prompt_embeds = text_encoder( + uncond_input.input_ids.to(device), + output_hidden_states=True, + ) + # We are only ALWAYS interested in the pooled output of the final text encoder + negative_pooled_prompt_embeds = negative_prompt_embeds[0] + negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2] + + negative_prompt_embeds_list.append(negative_prompt_embeds) + + negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1) + + prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device) + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( + bs_embed * num_images_per_prompt, -1 + ) + if do_classifier_free_guidance: + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view( + bs_embed * num_images_per_prompt, -1 + ) + + return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs( + self, + prompt, + prompt_2, + strength, + num_inference_steps, + callback_steps, + negative_prompt=None, + negative_prompt_2=None, + prompt_embeds=None, + negative_prompt_embeds=None, + ): + if strength < 0 or strength > 1: + raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}") + if num_inference_steps is None: + raise ValueError("`num_inference_steps` cannot be None.") + elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0: + raise ValueError( + f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type" + f" {type(num_inference_steps)}." + ) + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt_2 is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): + raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_2 is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + + def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None): + # get the original timestep using init_timestep + if denoising_start is None: + init_timestep = min(int(num_inference_steps * strength), num_inference_steps) + t_start = max(num_inference_steps - init_timestep, 0) + else: + t_start = 0 + + timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] + + # Strength is irrelevant if we directly request a timestep to start at; + # that is, strength is determined by the denoising_start instead. + if denoising_start is not None: + discrete_timestep_cutoff = int( + round( + self.scheduler.config.num_train_timesteps + - (denoising_start * self.scheduler.config.num_train_timesteps) + ) + ) + timesteps = list(filter(lambda ts: ts < discrete_timestep_cutoff, timesteps)) + return torch.tensor(timesteps), len(timesteps) + + return timesteps, num_inference_steps - t_start + + def prepare_latents( + self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True + ): + if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): + raise ValueError( + f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" + ) + + # Offload text encoder if `enable_model_cpu_offload` was enabled + if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: + self.text_encoder_2.to("cpu") + torch.cuda.empty_cache() + + image = image.to(device=device, dtype=dtype) + + batch_size = batch_size * num_images_per_prompt + + if image.shape[1] == 4: + init_latents = image + + else: + # make sure the VAE is in float32 mode, as it overflows in float16 + if self.vae.config.force_upcast: + image = image.float() + self.vae.to(dtype=torch.float32) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + elif isinstance(generator, list): + init_latents = [ + self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size) + ] + init_latents = torch.cat(init_latents, dim=0) + else: + init_latents = self.vae.encode(image).latent_dist.sample(generator) + + if self.vae.config.force_upcast: + self.vae.to(dtype) + + init_latents = init_latents.to(dtype) + init_latents = self.vae.config.scaling_factor * init_latents + + if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: + # expand init_latents for batch_size + additional_image_per_prompt = batch_size // init_latents.shape[0] + init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) + elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." + ) + else: + init_latents = torch.cat([init_latents], dim=0) + + if add_noise: + shape = init_latents.shape + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + # get latents + init_latents = self.scheduler.add_noise(init_latents, noise, timestep) + + latents = init_latents + + return latents + + def _get_add_time_ids( + self, + original_size, + crops_coords_top_left, + target_size, + aesthetic_score, + negative_aesthetic_score, + negative_original_size, + negative_crops_coords_top_left, + negative_target_size, + dtype, + ): + if self.config.requires_aesthetics_score: + add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) + add_neg_time_ids = list( + negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,) + ) + else: + add_time_ids = list(original_size + crops_coords_top_left + target_size) + add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size) + + passed_add_embed_dim = ( + self.unet.config.addition_time_embed_dim * len(add_time_ids) + self.text_encoder_2.config.projection_dim + ) + expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features + + if ( + expected_add_embed_dim > passed_add_embed_dim + and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_embed_dim + ): + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model." + ) + elif ( + expected_add_embed_dim < passed_add_embed_dim + and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_embed_dim + ): + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model." + ) + elif expected_add_embed_dim != passed_add_embed_dim: + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." + ) + + add_time_ids = torch.tensor([add_time_ids], dtype=dtype) + add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype) + + return add_time_ids, add_neg_time_ids + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae + def upcast_vae(self): + dtype = self.vae.dtype + self.vae.to(dtype=torch.float32) + use_torch_2_0_or_xformers = isinstance( + self.vae.decoder.mid_block.attentions[0].processor, + ( + AttnProcessor2_0, + XFormersAttnProcessor, + LoRAXFormersAttnProcessor, + LoRAAttnProcessor2_0, + ), + ) + # if xformers or torch_2_0 is used attention block does not need + # to be in float32 which can save lots of memory + if use_torch_2_0_or_xformers: + self.vae.post_quant_conv.to(dtype) + self.vae.decoder.conv_in.to(dtype) + self.vae.decoder.mid_block.to(dtype) + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + prompt_2: Optional[Union[str, List[str]]] = None, + image: PipelineImageInput = None, + strength: float = 0.3, + num_inference_steps: int = 50, + denoising_start: Optional[float] = None, + denoising_end: Optional[float] = None, + guidance_scale: float = 5.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt_2: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: int = 1, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guidance_rescale: float = 0.0, + original_size: Tuple[int, int] = None, + crops_coords_top_left: Tuple[int, int] = (0, 0), + target_size: Tuple[int, int] = None, + negative_original_size: Optional[Tuple[int, int]] = None, + negative_crops_coords_top_left: Tuple[int, int] = (0, 0), + negative_target_size: Optional[Tuple[int, int]] = None, + aesthetic_score: float = 6.0, + negative_aesthetic_score: float = 2.5, + cache_interval: int = 1, + cache_layer_id: int = None, + cache_block_id: int = None, + uniform: bool = True, + pow: float = None, + center: int = None, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + used in both text-encoders + image (`torch.FloatTensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.FloatTensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`): + The image(s) to modify with the pipeline. + strength (`float`, *optional*, defaults to 0.3): + Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image` + will be used as a starting point, adding more noise to it the larger the `strength`. The number of + denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will + be maximum and the denoising process will run for the full number of iterations specified in + `num_inference_steps`. A value of 1, therefore, essentially ignores `image`. Note that in the case of + `denoising_start` being declared as an integer, the value of `strength` will be ignored. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + denoising_start (`float`, *optional*): + When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be + bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and + it is assumed that the passed `image` is a partly denoised image. Note that when this is specified, + strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline + is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refining the Image + Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). + denoising_end (`float`, *optional*): + When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be + completed before it is intentionally prematurely terminated. As a result, the returned sample will + still retain a substantial amount of noise (ca. final 20% of timesteps still needed) and should be + denoised by a successor pipeline that has `denoising_start` set to 0.8 so that it only denoises the + final 20% of the scheduler. The denoising_end parameter should ideally be utilized when this pipeline + forms a part of a "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image + Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output). + guidance_scale (`float`, *optional*, defaults to 7.5): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + negative_prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and + `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to + [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `prompt` input argument. + negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` + input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] instead of a + plain tuple. + callback (`Callable`, *optional*): + A function that will be called every `callback_steps` steps during inference. The function will be + called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function will be called. If not specified, the callback will be + called at every step. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + guidance_rescale (`float`, *optional*, defaults to 0.7): + Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of + [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). + Guidance rescale factor should fix overexposure when using zero terminal SNR. + original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. + `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as + explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position + `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting + `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + For most cases, `target_size` should be set to the desired height and width of the generated image. If + not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in + section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + To negatively condition the generation process based on a specific image resolution. Part of SDXL's + micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more + information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. + negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): + To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's + micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more + information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. + negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): + To negatively condition the generation process based on a target image resolution. It should be as same + as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more + information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. + aesthetic_score (`float`, *optional*, defaults to 6.0): + Used to simulate an aesthetic score of the generated image by influencing the positive text condition. + Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). + negative_aesthetic_score (`float`, *optional*, defaults to 2.5): + Part of SDXL's micro-conditioning as explained in section 2.2 of + [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). Can be used to + simulate an aesthetic score of the generated image by influencing the negative text condition. + + Examples: + + Returns: + [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`: + [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a + `tuple. When returning a tuple, the first element is a list with the generated images. + """ + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + prompt_2, + strength, + num_inference_steps, + callback_steps, + negative_prompt, + negative_prompt_2, + prompt_embeds, + negative_prompt_embeds, + ) + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + text_encoder_lora_scale = ( + cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None + ) + ( + prompt_embeds, + negative_prompt_embeds, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + device=device, + num_images_per_prompt=num_images_per_prompt, + do_classifier_free_guidance=do_classifier_free_guidance, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + lora_scale=text_encoder_lora_scale, + ) + + # 4. Preprocess image + image = self.image_processor.preprocess(image) + + # 5. Prepare timesteps + def denoising_value_valid(dnv): + return isinstance(denoising_end, float) and 0 < dnv < 1 + + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps, num_inference_steps = self.get_timesteps( + num_inference_steps, strength, device, denoising_start=denoising_start if denoising_value_valid else None + ) + latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) + + add_noise = True if denoising_start is None else False + # 6. Prepare latent variables + latents = self.prepare_latents( + image, + latent_timestep, + batch_size, + num_images_per_prompt, + prompt_embeds.dtype, + device, + generator, + add_noise, + ) + # 7. Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + height, width = latents.shape[-2:] + height = height * self.vae_scale_factor + width = width * self.vae_scale_factor + + original_size = original_size or (height, width) + target_size = target_size or (height, width) + + # 8. Prepare added time ids & embeddings + if negative_original_size is None: + negative_original_size = original_size + if negative_target_size is None: + negative_target_size = target_size + + add_text_embeds = pooled_prompt_embeds + add_time_ids, add_neg_time_ids = self._get_add_time_ids( + original_size, + crops_coords_top_left, + target_size, + aesthetic_score, + negative_aesthetic_score, + negative_original_size, + negative_crops_coords_top_left, + negative_target_size, + dtype=prompt_embeds.dtype, + ) + add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1) + + if do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) + add_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1) + add_time_ids = torch.cat([add_neg_time_ids, add_time_ids], dim=0) + + prompt_embeds = prompt_embeds.to(device) + add_text_embeds = add_text_embeds.to(device) + add_time_ids = add_time_ids.to(device) + + # 9. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + + # 9.1 Apply denoising_end + if ( + denoising_end is not None + and denoising_start is not None + and denoising_value_valid(denoising_end) + and denoising_value_valid(denoising_start) + and denoising_start >= denoising_end + ): + raise ValueError( + f"`denoising_start`: {denoising_start} cannot be larger than or equal to `denoising_end`: " + + f" {denoising_end} when using type float." + ) + elif denoising_end is not None and denoising_value_valid(denoising_end): + discrete_timestep_cutoff = int( + round( + self.scheduler.config.num_train_timesteps + - (denoising_end * self.scheduler.config.num_train_timesteps) + ) + ) + num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) + timesteps = timesteps[:num_inference_steps] + + if cache_interval == 1: + interval_seq = list(range(num_inference_steps)) + else: + if uniform: + interval_seq = list(range(0, num_inference_steps, cache_interval)) + else: + num_slow_step = num_inference_steps//cache_interval + if num_inference_steps%cache_interval != 0: + num_slow_step += 1 + + interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,] + #print(interval_seq) + + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + if i in interval_seq: + prv_features = None + + # predict the noise residual + added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} + noise_pred, prv_features = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + added_cond_kwargs=added_cond_kwargs, + replicate_prv_feature=prv_features, + quick_replicate= cache_interval>1, + cache_layer_id=cache_layer_id, + cache_block_id=cache_block_id, + return_dict=False, + ) + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + if do_classifier_free_guidance and guidance_rescale > 0.0: + # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf + noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + callback(i, t, latents) + + if not output_type == "latent": + # make sure the VAE is in float32 mode, as it overflows in float16 + needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast + + if needs_upcasting: + self.upcast_vae() + latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype) + + image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] + + # cast back to fp16 if needed + if needs_upcasting: + self.vae.to(dtype=torch.float16) + else: + image = latents + return StableDiffusionXLPipelineOutput(images=image) + + # apply watermark if available + if self.watermark is not None: + image = self.watermark.apply_watermark(image) + + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return StableDiffusionXLPipelineOutput(images=image) + + # Overrride to properly handle the loading and unloading of the additional text encoder. + # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.load_lora_weights + def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs): + # We could have accessed the unet config from `lora_state_dict()` too. We pass + # it here explicitly to be able to tell that it's coming from an SDXL + # pipeline. + + # Remove any existing hooks. + if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"): + from accelerate.hooks import AlignDevicesHook, CpuOffload, remove_hook_from_module + else: + raise ImportError("Offloading requires `accelerate v0.17.0` or higher.") + + is_model_cpu_offload = False + is_sequential_cpu_offload = False + recursive = False + for _, component in self.components.items(): + if isinstance(component, torch.nn.Module): + if hasattr(component, "_hf_hook"): + is_model_cpu_offload = isinstance(getattr(component, "_hf_hook"), CpuOffload) + is_sequential_cpu_offload = isinstance(getattr(component, "_hf_hook"), AlignDevicesHook) + logger.info( + "Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again." + ) + recursive = is_sequential_cpu_offload + remove_hook_from_module(component, recurse=recursive) + state_dict, network_alphas = self.lora_state_dict( + pretrained_model_name_or_path_or_dict, + unet_config=self.unet.config, + **kwargs, + ) + self.load_lora_into_unet(state_dict, network_alphas=network_alphas, unet=self.unet) + + text_encoder_state_dict = {k: v for k, v in state_dict.items() if "text_encoder." in k} + if len(text_encoder_state_dict) > 0: + self.load_lora_into_text_encoder( + text_encoder_state_dict, + network_alphas=network_alphas, + text_encoder=self.text_encoder, + prefix="text_encoder", + lora_scale=self.lora_scale, + ) + + text_encoder_2_state_dict = {k: v for k, v in state_dict.items() if "text_encoder_2." in k} + if len(text_encoder_2_state_dict) > 0: + self.load_lora_into_text_encoder( + text_encoder_2_state_dict, + network_alphas=network_alphas, + text_encoder=self.text_encoder_2, + prefix="text_encoder_2", + lora_scale=self.lora_scale, + ) + + # Offload back. + if is_model_cpu_offload: + self.enable_model_cpu_offload() + elif is_sequential_cpu_offload: + self.enable_sequential_cpu_offload() + + @classmethod + # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.save_lora_weights + def save_lora_weights( + self, + save_directory: Union[str, os.PathLike], + unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None, + is_main_process: bool = True, + weight_name: str = None, + save_function: Callable = None, + safe_serialization: bool = True, + ): + state_dict = {} + + def pack_weights(layers, prefix): + layers_weights = layers.state_dict() if isinstance(layers, torch.nn.Module) else layers + layers_state_dict = {f"{prefix}.{module_name}": param for module_name, param in layers_weights.items()} + return layers_state_dict + + if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers): + raise ValueError( + "You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers` or `text_encoder_2_lora_layers`." + ) + + if unet_lora_layers: + state_dict.update(pack_weights(unet_lora_layers, "unet")) + + if text_encoder_lora_layers and text_encoder_2_lora_layers: + state_dict.update(pack_weights(text_encoder_lora_layers, "text_encoder")) + state_dict.update(pack_weights(text_encoder_2_lora_layers, "text_encoder_2")) + + self.write_lora_layers( + state_dict=state_dict, + save_directory=save_directory, + is_main_process=is_main_process, + weight_name=weight_name, + save_function=save_function, + safe_serialization=safe_serialization, + ) + + # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline._remove_text_encoder_monkey_patch + def _remove_text_encoder_monkey_patch(self): + self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder) + self._remove_text_encoder_monkey_patch_classmethod(self.text_encoder_2) diff --git a/ixformer_sdk/contrib/DeepCache/sdxl/pipeline_utils.py b/ixformer_sdk/contrib/DeepCache/sdxl/pipeline_utils.py new file mode 100644 index 00000000..41c1c2a7 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sdxl/pipeline_utils.py @@ -0,0 +1,1839 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# +# 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 fnmatch +import importlib +import inspect +import os +import re +import sys +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import PIL +import torch +from huggingface_hub import ModelCard, create_repo, hf_hub_download, model_info, snapshot_download +from packaging import version +from requests.exceptions import HTTPError +from tqdm.auto import tqdm + +import diffusers + +from diffusers import __version__ +from diffusers.configuration_utils import ConfigMixin +from diffusers.models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT +from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME +from diffusers.utils import ( + CONFIG_NAME, + DEPRECATED_REVISION_ARGS, + # DIFFUSERS_CACHE, + # HF_HUB_OFFLINE, + SAFETENSORS_WEIGHTS_NAME, + WEIGHTS_NAME, + BaseOutput, + deprecate, + get_class_from_dynamic_module, + is_accelerate_available, + is_accelerate_version, + is_torch_version, + is_transformers_available, + logging, + numpy_to_pil, +) +from diffusers.utils.torch_utils import is_compiled_module +from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE +DIFFUSERS_CACHE=HUGGINGFACE_HUB_CACHE +ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} +HF_HUB_OFFLINE = os.getenv("HF_HUB_OFFLINE", "").upper() in ENV_VARS_TRUE_VALUES +if is_transformers_available(): + import transformers + from transformers import PreTrainedModel + from transformers.utils import FLAX_WEIGHTS_NAME as TRANSFORMERS_FLAX_WEIGHTS_NAME + from transformers.utils import SAFE_WEIGHTS_NAME as TRANSFORMERS_SAFE_WEIGHTS_NAME + from transformers.utils import WEIGHTS_NAME as TRANSFORMERS_WEIGHTS_NAME + +from diffusers.utils import FLAX_WEIGHTS_NAME, ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, PushToHubMixin + + +if is_accelerate_available(): + import accelerate + + +INDEX_FILE = "diffusion_pytorch_model.bin" +CUSTOM_PIPELINE_FILE_NAME = "pipeline.py" +DUMMY_MODULES_FOLDER = "diffusers.utils" +TRANSFORMERS_DUMMY_MODULES_FOLDER = "transformers.utils" +CONNECTED_PIPES_KEYS = ["prior"] + + +logger = logging.get_logger(__name__) + + +LOADABLE_CLASSES = { + "diffusers": { + "ModelMixin": ["save_pretrained", "from_pretrained"], + "SchedulerMixin": ["save_pretrained", "from_pretrained"], + "DiffusionPipeline": ["save_pretrained", "from_pretrained"], + "OnnxRuntimeModel": ["save_pretrained", "from_pretrained"], + }, + "transformers": { + "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"], + "PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"], + "PreTrainedModel": ["save_pretrained", "from_pretrained"], + "FeatureExtractionMixin": ["save_pretrained", "from_pretrained"], + "ProcessorMixin": ["save_pretrained", "from_pretrained"], + "ImageProcessingMixin": ["save_pretrained", "from_pretrained"], + }, + "onnxruntime.training": { + "ORTModule": ["save_pretrained", "from_pretrained"], + }, +} + +ALL_IMPORTABLE_CLASSES = {} +for library in LOADABLE_CLASSES: + ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library]) + + +@dataclass +class ImagePipelineOutput(BaseOutput): + """ + Output class for image pipelines. + + Args: + images (`List[PIL.Image.Image]` or `np.ndarray`) + List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + num_channels)`. + """ + + images: Union[List[PIL.Image.Image], np.ndarray] + + +@dataclass +class AudioPipelineOutput(BaseOutput): + """ + Output class for audio pipelines. + + Args: + audios (`np.ndarray`) + List of denoised audio samples of a NumPy array of shape `(batch_size, num_channels, sample_rate)`. + """ + + audios: np.ndarray + + +def is_safetensors_compatible(filenames, variant=None, passed_components=None) -> bool: + """ + Checking for safetensors compatibility: + - By default, all models are saved with the default pytorch serialization, so we use the list of default pytorch + files to know which safetensors files are needed. + - The model is safetensors compatible only if there is a matching safetensors file for every default pytorch file. + + Converting default pytorch serialized filenames to safetensors serialized filenames: + - For models from the diffusers library, just replace the ".bin" extension with ".safetensors" + - For models from the transformers library, the filename changes from "pytorch_model" to "model", and the ".bin" + extension is replaced with ".safetensors" + """ + pt_filenames = [] + + sf_filenames = set() + + passed_components = passed_components or [] + + for filename in filenames: + _, extension = os.path.splitext(filename) + + if len(filename.split("/")) == 2 and filename.split("/")[0] in passed_components: + continue + + if extension == ".bin": + pt_filenames.append(filename) + elif extension == ".safetensors": + sf_filenames.add(filename) + + for filename in pt_filenames: + # filename = 'foo/bar/baz.bam' -> path = 'foo/bar', filename = 'baz', extention = '.bam' + path, filename = os.path.split(filename) + filename, extension = os.path.splitext(filename) + + if filename.startswith("pytorch_model"): + filename = filename.replace("pytorch_model", "model") + else: + filename = filename + + expected_sf_filename = os.path.join(path, filename) + expected_sf_filename = f"{expected_sf_filename}.safetensors" + + if expected_sf_filename not in sf_filenames: + logger.warning(f"{expected_sf_filename} not found") + return False + + return True + + +def variant_compatible_siblings(filenames, variant=None) -> Union[List[os.PathLike], str]: + weight_names = [ + WEIGHTS_NAME, + SAFETENSORS_WEIGHTS_NAME, + FLAX_WEIGHTS_NAME, + ONNX_WEIGHTS_NAME, + ONNX_EXTERNAL_WEIGHTS_NAME, + ] + + if is_transformers_available(): + weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME, TRANSFORMERS_FLAX_WEIGHTS_NAME] + + # model_pytorch, diffusion_model_pytorch, ... + weight_prefixes = [w.split(".")[0] for w in weight_names] + # .bin, .safetensors, ... + weight_suffixs = [w.split(".")[-1] for w in weight_names] + # -00001-of-00002 + transformers_index_format = r"\d{5}-of-\d{5}" + + if variant is not None: + # `diffusion_pytorch_model.fp16.bin` as well as `model.fp16-00001-of-00002.safetensors` + variant_file_re = re.compile( + rf"({'|'.join(weight_prefixes)})\.({variant}|{variant}-{transformers_index_format})\.({'|'.join(weight_suffixs)})$" + ) + # `text_encoder/pytorch_model.bin.index.fp16.json` + variant_index_re = re.compile( + rf"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.{variant}\.json$" + ) + + # `diffusion_pytorch_model.bin` as well as `model-00001-of-00002.safetensors` + non_variant_file_re = re.compile( + rf"({'|'.join(weight_prefixes)})(-{transformers_index_format})?\.({'|'.join(weight_suffixs)})$" + ) + # `text_encoder/pytorch_model.bin.index.json` + non_variant_index_re = re.compile(rf"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.json") + + if variant is not None: + variant_weights = {f for f in filenames if variant_file_re.match(f.split("/")[-1]) is not None} + variant_indexes = {f for f in filenames if variant_index_re.match(f.split("/")[-1]) is not None} + variant_filenames = variant_weights | variant_indexes + else: + variant_filenames = set() + + non_variant_weights = {f for f in filenames if non_variant_file_re.match(f.split("/")[-1]) is not None} + non_variant_indexes = {f for f in filenames if non_variant_index_re.match(f.split("/")[-1]) is not None} + non_variant_filenames = non_variant_weights | non_variant_indexes + + # all variant filenames will be used by default + usable_filenames = set(variant_filenames) + + def convert_to_variant(filename): + if "index" in filename: + variant_filename = filename.replace("index", f"index.{variant}") + elif re.compile(f"^(.*?){transformers_index_format}").match(filename) is not None: + variant_filename = f"{filename.split('-')[0]}.{variant}-{'-'.join(filename.split('-')[1:])}" + else: + variant_filename = f"{filename.split('.')[0]}.{variant}.{filename.split('.')[1]}" + return variant_filename + + for f in non_variant_filenames: + variant_filename = convert_to_variant(f) + if variant_filename not in usable_filenames: + usable_filenames.add(f) + + return usable_filenames, variant_filenames + + +def warn_deprecated_model_variant(pretrained_model_name_or_path, use_auth_token, variant, revision, model_filenames): + info = model_info( + pretrained_model_name_or_path, + use_auth_token=use_auth_token, + revision=None, + ) + filenames = {sibling.rfilename for sibling in info.siblings} + comp_model_filenames, _ = variant_compatible_siblings(filenames, variant=revision) + comp_model_filenames = [".".join(f.split(".")[:1] + f.split(".")[2:]) for f in comp_model_filenames] + + if set(comp_model_filenames) == set(model_filenames): + warnings.warn( + f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'` even though you can load it via `variant=`{revision}`. Loading model variants via `revision='{revision}'` is deprecated and will be removed in diffusers v1. Please use `variant='{revision}'` instead.", + FutureWarning, + ) + else: + warnings.warn( + f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant='{revision}'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have the required variant filenames in the 'main' branch. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title '{pretrained_model_name_or_path} is missing {revision} files' so that the correct variant file can be added.", + FutureWarning, + ) + + +def maybe_raise_or_warn( + library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module +): + """Simple helper method to raise or warn in case incorrect module has been passed""" + if not is_pipeline_module: + library = importlib.import_module(library_name) + class_obj = getattr(library, class_name) + class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} + + expected_class_obj = None + for class_name, class_candidate in class_candidates.items(): + if class_candidate is not None and issubclass(class_obj, class_candidate): + expected_class_obj = class_candidate + + # Dynamo wraps the original model in a private class. + # I didn't find a public API to get the original class. + sub_model = passed_class_obj[name] + model_cls = sub_model.__class__ + if is_compiled_module(sub_model): + model_cls = sub_model._orig_mod.__class__ + + if not issubclass(model_cls, expected_class_obj): + raise ValueError( + f"{passed_class_obj[name]} is of type: {model_cls}, but should be" f" {expected_class_obj}" + ) + else: + logger.warning( + f"You have passed a non-standard module {passed_class_obj[name]}. We cannot verify whether it" + " has the correct type" + ) + + +def get_class_obj_and_candidates(library_name, class_name, importable_classes, pipelines, is_pipeline_module): + """Simple helper method to retrieve class object of module as well as potential parent class objects""" + if is_pipeline_module: + pipeline_module = getattr(pipelines, library_name) + + class_obj = getattr(pipeline_module, class_name) + class_candidates = {c: class_obj for c in importable_classes.keys()} + else: + # else we just import it from the library. + if class_name == 'UNet2DConditionModel': + library_name = "ixformer.contrib.DeepCache.sdxl.unet_2d_condition" + + library = importlib.import_module(library_name) + class_obj = getattr(library, class_name) + class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} + + return class_obj, class_candidates + + +def _get_pipeline_class( + class_obj, config, load_connected_pipeline=False, custom_pipeline=None, cache_dir=None, revision=None +): + if custom_pipeline is not None: + if custom_pipeline.endswith(".py"): + path = Path(custom_pipeline) + # decompose into folder & file + file_name = path.name + custom_pipeline = path.parent.absolute() + else: + file_name = CUSTOM_PIPELINE_FILE_NAME + + return get_class_from_dynamic_module( + custom_pipeline, module_file=file_name, cache_dir=cache_dir, revision=revision + ) + + if class_obj != DiffusionPipeline: + return class_obj + + diffusers_module = importlib.import_module(class_obj.__module__.split(".")[0]) + class_name = config["_class_name"] + + if class_name.startswith("Flax"): + class_name = class_name[4:] + + pipeline_cls = getattr(diffusers_module, class_name) + + if load_connected_pipeline: + from .auto_pipeline import _get_connected_pipeline + + connected_pipeline_cls = _get_connected_pipeline(pipeline_cls) + if connected_pipeline_cls is not None: + logger.info( + f"Loading connected pipeline {connected_pipeline_cls.__name__} instead of {pipeline_cls.__name__} as specified via `load_connected_pipeline=True`" + ) + else: + logger.info(f"{pipeline_cls.__name__} has no connected pipeline class. Loading {pipeline_cls.__name__}.") + + pipeline_cls = connected_pipeline_cls or pipeline_cls + + return pipeline_cls + + +def load_sub_model( + library_name: str, + class_name: str, + importable_classes: List[Any], + pipelines: Any, + is_pipeline_module: bool, + pipeline_class: Any, + torch_dtype: torch.dtype, + provider: Any, + sess_options: Any, + device_map: Optional[Union[Dict[str, torch.device], str]], + max_memory: Optional[Dict[Union[int, str], Union[int, str]]], + offload_folder: Optional[Union[str, os.PathLike]], + offload_state_dict: bool, + model_variants: Dict[str, str], + name: str, + from_flax: bool, + variant: str, + low_cpu_mem_usage: bool, + cached_folder: Union[str, os.PathLike], +): + """Helper method to load the module `name` from `library_name` and `class_name`""" + # retrieve class candidates + class_obj, class_candidates = get_class_obj_and_candidates( + library_name, class_name, importable_classes, pipelines, is_pipeline_module + ) + + load_method_name = None + # retrive load method name + for class_name, class_candidate in class_candidates.items(): + if class_candidate is not None and issubclass(class_obj, class_candidate): + load_method_name = importable_classes[class_name][1] + + # if load method name is None, then we have a dummy module -> raise Error + if load_method_name is None: + none_module = class_obj.__module__ + is_dummy_path = none_module.startswith(DUMMY_MODULES_FOLDER) or none_module.startswith( + TRANSFORMERS_DUMMY_MODULES_FOLDER + ) + if is_dummy_path and "dummy" in none_module: + # call class_obj for nice error message of missing requirements + class_obj() + + raise ValueError( + f"The component {class_obj} of {pipeline_class} cannot be loaded as it does not seem to have" + f" any of the loading methods defined in {ALL_IMPORTABLE_CLASSES}." + ) + + load_method = getattr(class_obj, load_method_name) + + # add kwargs to loading method + loading_kwargs = {} + if issubclass(class_obj, torch.nn.Module): + loading_kwargs["torch_dtype"] = torch_dtype + if issubclass(class_obj, diffusers.OnnxRuntimeModel): + loading_kwargs["provider"] = provider + loading_kwargs["sess_options"] = sess_options + + is_diffusers_model = issubclass(class_obj, diffusers.ModelMixin) + + if is_transformers_available(): + transformers_version = version.parse(version.parse(transformers.__version__).base_version) + else: + transformers_version = "N/A" + + is_transformers_model = ( + is_transformers_available() + and issubclass(class_obj, PreTrainedModel) + and transformers_version >= version.parse("4.20.0") + ) + + # When loading a transformers model, if the device_map is None, the weights will be initialized as opposed to diffusers. + # To make default loading faster we set the `low_cpu_mem_usage=low_cpu_mem_usage` flag which is `True` by default. + # This makes sure that the weights won't be initialized which significantly speeds up loading. + if is_diffusers_model or is_transformers_model: + loading_kwargs["device_map"] = device_map + loading_kwargs["max_memory"] = max_memory + loading_kwargs["offload_folder"] = offload_folder + loading_kwargs["offload_state_dict"] = offload_state_dict + loading_kwargs["variant"] = model_variants.pop(name, None) + if from_flax: + loading_kwargs["from_flax"] = True + + # the following can be deleted once the minimum required `transformers` version + # is higher than 4.27 + if ( + is_transformers_model + and loading_kwargs["variant"] is not None + and transformers_version < version.parse("4.27.0") + ): + raise ImportError( + f"When passing `variant='{variant}'`, please make sure to upgrade your `transformers` version to at least 4.27.0.dev0" + ) + elif is_transformers_model and loading_kwargs["variant"] is None: + loading_kwargs.pop("variant") + + # if `from_flax` and model is transformer model, can currently not load with `low_cpu_mem_usage` + if not (from_flax and is_transformers_model): + loading_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage + else: + loading_kwargs["low_cpu_mem_usage"] = False + + # check if the module is in a subdirectory + if os.path.isdir(os.path.join(cached_folder, name)): + loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs) + else: + # else load from the root directory + loaded_sub_model = load_method(cached_folder, **loading_kwargs) + + return loaded_sub_model + + +class DiffusionPipeline(ConfigMixin, PushToHubMixin): + r""" + Base class for all pipelines. + + [`DiffusionPipeline`] stores all components (models, schedulers, and processors) for diffusion pipelines and + provides methods for loading, downloading and saving models. It also includes methods to: + + - move all PyTorch modules to the device of your choice + - enable/disable the progress bar for the denoising iteration + + Class attributes: + + - **config_name** (`str`) -- The configuration filename that stores the class and module names of all the + diffusion pipeline's components. + - **_optional_components** (`List[str]`) -- List of all optional components that don't have to be passed to the + pipeline to function (should be overridden by subclasses). + """ + config_name = "model_index.json" + model_cpu_offload_seq = None + _optional_components = [] + _exclude_from_cpu_offload = [] + _load_connected_pipes = False + _is_onnx = False + + def register_modules(self, **kwargs): + # import it here to avoid circular import + from diffusers import pipelines + + for name, module in kwargs.items(): + # retrieve library + if module is None: + register_dict = {name: (None, None)} + else: + # register the config from the original module, not the dynamo compiled one + if is_compiled_module(module): + not_compiled_module = module._orig_mod + else: + not_compiled_module = module + + library = not_compiled_module.__module__.split(".")[0] + + # check if the module is a pipeline module + module_path_items = not_compiled_module.__module__.split(".") + pipeline_dir = module_path_items[-2] if len(module_path_items) > 2 else None + + path = not_compiled_module.__module__.split(".") + is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir) + + # if library is not in LOADABLE_CLASSES, then it is a custom module. + # Or if it's a pipeline module, then the module is inside the pipeline + # folder so we set the library to module name. + if is_pipeline_module: + library = pipeline_dir + elif library not in LOADABLE_CLASSES: + library = not_compiled_module.__module__ + + # retrieve class_name + class_name = not_compiled_module.__class__.__name__ + + register_dict = {name: (library, class_name)} + + # save model index config + self.register_to_config(**register_dict) + + # set models + setattr(self, name, module) + + def __setattr__(self, name: str, value: Any): + if name in self.__dict__ and hasattr(self.config, name): + # We need to overwrite the config if name exists in config + if isinstance(getattr(self.config, name), (tuple, list)): + if value is not None and self.config[name][0] is not None: + class_library_tuple = (value.__module__.split(".")[0], value.__class__.__name__) + else: + class_library_tuple = (None, None) + + self.register_to_config(**{name: class_library_tuple}) + else: + self.register_to_config(**{name: value}) + + super().__setattr__(name, value) + + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + safe_serialization: bool = True, + variant: Optional[str] = None, + push_to_hub: bool = False, + **kwargs, + ): + """ + Save all saveable variables of the pipeline to a directory. A pipeline variable can be saved and loaded if its + class implements both a save and loading method. The pipeline is easily reloaded using the + [`~DiffusionPipeline.from_pretrained`] class method. + + Arguments: + save_directory (`str` or `os.PathLike`): + Directory to save a pipeline to. Will be created if it doesn't exist. + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. + variant (`str`, *optional*): + If specified, weights are saved in the format `pytorch_model..bin`. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + model_index_dict = dict(self.config) + model_index_dict.pop("_class_name", None) + model_index_dict.pop("_diffusers_version", None) + model_index_dict.pop("_module", None) + model_index_dict.pop("_name_or_path", None) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + private = kwargs.pop("private", False) + create_pr = kwargs.pop("create_pr", False) + token = kwargs.pop("token", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id + + expected_modules, optional_kwargs = self._get_signature_keys(self) + + def is_saveable_module(name, value): + if name not in expected_modules: + return False + if name in self._optional_components and value[0] is None: + return False + return True + + model_index_dict = {k: v for k, v in model_index_dict.items() if is_saveable_module(k, v)} + for pipeline_component_name in model_index_dict.keys(): + sub_model = getattr(self, pipeline_component_name) + model_cls = sub_model.__class__ + + # Dynamo wraps the original model in a private class. + # I didn't find a public API to get the original class. + if is_compiled_module(sub_model): + sub_model = sub_model._orig_mod + model_cls = sub_model.__class__ + + save_method_name = None + # search for the model's base class in LOADABLE_CLASSES + for library_name, library_classes in LOADABLE_CLASSES.items(): + if library_name in sys.modules: + library = importlib.import_module(library_name) + else: + logger.info( + f"{library_name} is not installed. Cannot save {pipeline_component_name} as {library_classes} from {library_name}" + ) + + for base_class, save_load_methods in library_classes.items(): + class_candidate = getattr(library, base_class, None) + if class_candidate is not None and issubclass(model_cls, class_candidate): + # if we found a suitable base class in LOADABLE_CLASSES then grab its save method + save_method_name = save_load_methods[0] + break + if save_method_name is not None: + break + + if save_method_name is None: + logger.warn(f"self.{pipeline_component_name}={sub_model} of type {type(sub_model)} cannot be saved.") + # make sure that unsaveable components are not tried to be loaded afterward + self.register_to_config(**{pipeline_component_name: (None, None)}) + continue + + save_method = getattr(sub_model, save_method_name) + + # Call the save method with the argument safe_serialization only if it's supported + save_method_signature = inspect.signature(save_method) + save_method_accept_safe = "safe_serialization" in save_method_signature.parameters + save_method_accept_variant = "variant" in save_method_signature.parameters + + save_kwargs = {} + if save_method_accept_safe: + save_kwargs["safe_serialization"] = safe_serialization + if save_method_accept_variant: + save_kwargs["variant"] = variant + + save_method(os.path.join(save_directory, pipeline_component_name), **save_kwargs) + + # finally save the config + self.save_config(save_directory) + + if push_to_hub: + self._upload_folder( + save_directory, + repo_id, + token=token, + commit_message=commit_message, + create_pr=create_pr, + ) + + def to( + self, + torch_device: Optional[Union[str, torch.device]] = None, + torch_dtype: Optional[torch.dtype] = None, + silence_dtype_warnings: bool = False, + ): + if torch_device is None and torch_dtype is None: + return self + + # throw warning if pipeline is in "offloaded"-mode but user tries to manually set to GPU. + def module_is_sequentially_offloaded(module): + if not is_accelerate_available() or is_accelerate_version("<", "0.14.0"): + return False + + return hasattr(module, "_hf_hook") and not isinstance( + module._hf_hook, (accelerate.hooks.CpuOffload, accelerate.hooks.AlignDevicesHook) + ) + + def module_is_offloaded(module): + if not is_accelerate_available() or is_accelerate_version("<", "0.17.0.dev0"): + return False + + return hasattr(module, "_hf_hook") and isinstance(module._hf_hook, accelerate.hooks.CpuOffload) + + # .to("cuda") would raise an error if the pipeline is sequentially offloaded, so we raise our own to make it clearer + pipeline_is_sequentially_offloaded = any( + module_is_sequentially_offloaded(module) for _, module in self.components.items() + ) + if pipeline_is_sequentially_offloaded and torch_device and torch.device(torch_device).type == "cuda": + raise ValueError( + "It seems like you have activated sequential model offloading by calling `enable_sequential_cpu_offload`, but are now attempting to move the pipeline to GPU. This is not compatible with offloading. Please, move your pipeline `.to('cpu')` or consider removing the move altogether if you use sequential offloading." + ) + + # Display a warning in this case (the operation succeeds but the benefits are lost) + pipeline_is_offloaded = any(module_is_offloaded(module) for _, module in self.components.items()) + if pipeline_is_offloaded and torch_device and torch.device(torch_device).type == "cuda": + logger.warning( + f"It seems like you have activated model offloading by calling `enable_model_cpu_offload`, but are now manually moving the pipeline to GPU. It is strongly recommended against doing so as memory gains from offloading are likely to be lost. Offloading automatically takes care of moving the individual components {', '.join(self.components.keys())} to GPU when needed. To make sure offloading works as expected, you should consider moving the pipeline back to CPU: `pipeline.to('cpu')` or removing the move altogether if you use offloading." + ) + + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + is_offloaded = pipeline_is_offloaded or pipeline_is_sequentially_offloaded + for module in modules: + is_loaded_in_8bit = hasattr(module, "is_loaded_in_8bit") and module.is_loaded_in_8bit + + if is_loaded_in_8bit and torch_dtype is not None: + logger.warning( + f"The module '{module.__class__.__name__}' has been loaded in 8bit and conversion to {torch_dtype} is not yet supported. Module is still in 8bit precision." + ) + + if is_loaded_in_8bit and torch_device is not None: + logger.warning( + f"The module '{module.__class__.__name__}' has been loaded in 8bit and moving it to {torch_dtype} via `.to()` is not yet supported. Module is still on {module.device}." + ) + else: + module.to(torch_device, torch_dtype) + + if ( + module.dtype == torch.float16 + and str(torch_device) in ["cpu"] + and not silence_dtype_warnings + and not is_offloaded + ): + logger.warning( + "Pipelines loaded with `torch_dtype=torch.float16` cannot run with `cpu` device. It" + " is not recommended to move them to `cpu` as running them will fail. Please make" + " sure to use an accelerator to run the pipeline in inference, due to the lack of" + " support for`float16` operations on this device in PyTorch. Please, remove the" + " `torch_dtype=torch.float16` argument, or use another device for inference." + ) + return self + + @property + def device(self) -> torch.device: + r""" + Returns: + `torch.device`: The torch device on which the pipeline is located. + """ + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + for module in modules: + return module.device + + return torch.device("cpu") + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): + r""" + Instantiate a PyTorch diffusion pipeline from pretrained pipeline weights. + + The pipeline is set in evaluation mode (`model.eval()`) by default. + + If you get the error message below, you need to finetune the weights for your downstream task: + + ``` + Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: + - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated + You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. + ``` + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): + Can be either: + + - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline + hosted on the Hub. + - A path to a *directory* (for example `./my_pipeline_directory/`) containing pipeline weights + saved using + [`~DiffusionPipeline.save_pretrained`]. + torch_dtype (`str` or `torch.dtype`, *optional*): + Override the default `torch.dtype` and load the model with another dtype. If "auto" is passed, the + dtype is automatically derived from the model's weights. + custom_pipeline (`str`, *optional*): + + + + 🧪 This is an experimental feature and may change in the future. + + + + Can be either: + + - A string, the *repo id* (for example `hf-internal-testing/diffusers-dummy-pipeline`) of a custom + pipeline hosted on the Hub. The repository must contain a file called pipeline.py that defines + the custom pipeline. + - A string, the *file name* of a community pipeline hosted on GitHub under + [Community](https://github.com/huggingface/diffusers/tree/main/examples/community). Valid file + names must match the file name and not the pipeline script (`clip_guided_stable_diffusion` + instead of `clip_guided_stable_diffusion.py`). Community pipelines are always loaded from the + current main branch of GitHub. + - A path to a directory (`./my_pipeline_directory/`) containing a custom pipeline. The directory + must contain a file called `pipeline.py` that defines the custom pipeline. + + For more information on how to load and create custom pipelines, please have a look at [Loading and + Adding Custom + Pipelines](https://huggingface.co/docs/diffusers/using-diffusers/custom_pipeline_overview) + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory where a downloaded pretrained model configuration is cached if the standard cache + is not used. + resume_download (`bool`, *optional*, defaults to `False`): + Whether or not to resume downloading the model weights and configuration files. If set to `False`, any + incompletely downloaded files are deleted. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info(`bool`, *optional*, defaults to `False`): + Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files or not. If set to `True`, the model + won't be downloaded from the Hub. + use_auth_token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from + `diffusers-cli login` (stored in `~/.huggingface`) is used. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier + allowed by Git. + custom_revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id similar to + `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a + custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. + mirror (`str`, *optional*): + Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not + guarantee the timeliness or safety of the source, and you should refer to the mirror site for more + information. + device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + A map that specifies where each submodule should go. It doesn’t need to be defined for each + parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the + same device. + + Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For + more information about each option see [designing a device + map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). + max_memory (`Dict`, *optional*): + A dictionary device identifier for the maximum memory. Will default to the maximum memory available for + each GPU and the available CPU RAM if unset. + offload_folder (`str` or `os.PathLike`, *optional*): + The path to offload weights if device_map contains the value `"disk"`. + offload_state_dict (`bool`, *optional*): + If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if + the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` + when there is some disk offload. + low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): + Speed up model loading only loading the pretrained weights and not initializing the weights. This also + tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. + Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this + argument to `True` will raise an error. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `None`, the safetensors weights are downloaded if they're available **and** if the + safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors + weights. If set to `False`, safetensors weights are not loaded. + use_onnx (`bool`, *optional*, defaults to `None`): + If set to `True`, ONNX weights will always be downloaded if present. If set to `False`, ONNX weights + will never be downloaded. By default `use_onnx` defaults to the `_is_onnx` class attribute which is + `False` for non-ONNX pipelines and `True` for ONNX pipelines. ONNX weights include both files ending + with `.onnx` and `.pb`. + kwargs (remaining dictionary of keyword arguments, *optional*): + Can be used to overwrite load and saveable variables (the pipeline components of the specific pipeline + class). The overwritten components are passed directly to the pipelines `__init__` method. See example + below for more information. + variant (`str`, *optional*): + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when + loading `from_flax`. + + + + To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with + `huggingface-cli login`. + + + + Examples: + + ```py + >>> from diffusers import DiffusionPipeline + + >>> # Download pipeline from huggingface.co and cache. + >>> pipeline = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256") + + >>> # Download pipeline that requires an authorization token + >>> # For more information on access tokens, please refer to this section + >>> # of the documentation](https://huggingface.co/docs/hub/security-tokens) + >>> pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5") + + >>> # Use a different scheduler + >>> from diffusers import LMSDiscreteScheduler + + >>> scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config) + >>> pipeline.scheduler = scheduler + ``` + """ + cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE) + resume_download = kwargs.pop("resume_download", False) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + from_flax = kwargs.pop("from_flax", False) + torch_dtype = kwargs.pop("torch_dtype", None) + custom_pipeline = kwargs.pop("custom_pipeline", None) + custom_revision = kwargs.pop("custom_revision", None) + provider = kwargs.pop("provider", None) + sess_options = kwargs.pop("sess_options", None) + device_map = kwargs.pop("device_map", None) + max_memory = kwargs.pop("max_memory", None) + offload_folder = kwargs.pop("offload_folder", None) + offload_state_dict = kwargs.pop("offload_state_dict", False) + low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None) + use_onnx = kwargs.pop("use_onnx", None) + load_connected_pipeline = kwargs.pop("load_connected_pipeline", False) + + print("In our loading pipeline") + # 1. Download the checkpoints and configs + # use snapshot download here to get it working from from_pretrained + if not os.path.isdir(pretrained_model_name_or_path): + cached_folder = cls.download( + pretrained_model_name_or_path, + cache_dir=cache_dir, + resume_download=resume_download, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + from_flax=from_flax, + use_safetensors=use_safetensors, + use_onnx=use_onnx, + custom_pipeline=custom_pipeline, + custom_revision=custom_revision, + variant=variant, + load_connected_pipeline=load_connected_pipeline, + **kwargs, + ) + else: + cached_folder = pretrained_model_name_or_path + + config_dict = cls.load_config(cached_folder) + + # pop out "_ignore_files" as it is only needed for download + config_dict.pop("_ignore_files", None) + + # 2. Define which model components should load variants + # We retrieve the information by matching whether variant + # model checkpoints exist in the subfolders + model_variants = {} + if variant is not None: + for folder in os.listdir(cached_folder): + folder_path = os.path.join(cached_folder, folder) + is_folder = os.path.isdir(folder_path) and folder in config_dict + variant_exists = is_folder and any( + p.split(".")[1].startswith(variant) for p in os.listdir(folder_path) + ) + if variant_exists: + model_variants[folder] = variant + + # 3. Load the pipeline class, if using custom module then load it from the hub + # if we load from explicit class, let's use it + pipeline_class = _get_pipeline_class( + cls, + config_dict, + load_connected_pipeline=load_connected_pipeline, + custom_pipeline=custom_pipeline, + cache_dir=cache_dir, + revision=custom_revision, + ) + + # DEPRECATED: To be removed in 1.0.0 + if pipeline_class.__name__ == "StableDiffusionInpaintPipeline" and version.parse( + version.parse(config_dict["_diffusers_version"]).base_version + ) <= version.parse("0.5.1"): + from diffusers import StableDiffusionInpaintPipeline, StableDiffusionInpaintPipelineLegacy + + pipeline_class = StableDiffusionInpaintPipelineLegacy + + deprecation_message = ( + "You are using a legacy checkpoint for inpainting with Stable Diffusion, therefore we are loading the" + f" {StableDiffusionInpaintPipelineLegacy} class instead of {StableDiffusionInpaintPipeline}. For" + " better inpainting results, we strongly suggest using Stable Diffusion's official inpainting" + " checkpoint: https://huggingface.co/runwayml/stable-diffusion-inpainting instead or adapting your" + f" checkpoint {pretrained_model_name_or_path} to the format of" + " https://huggingface.co/runwayml/stable-diffusion-inpainting. Note that we do not actively maintain" + " the {StableDiffusionInpaintPipelineLegacy} class and will likely remove it in version 1.0.0." + ) + deprecate("StableDiffusionInpaintPipelineLegacy", "1.0.0", deprecation_message, standard_warn=False) + + # 4. Define expected modules given pipeline signature + # and define non-None initialized modules (=`init_kwargs`) + + # some modules can be passed directly to the init + # in this case they are already instantiated in `kwargs` + # extract them here + expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class) + passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} + passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} + + init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs) + + # define init kwargs and make sure that optional component modules are filtered out + init_kwargs = { + k: init_dict.pop(k) + for k in optional_kwargs + if k in init_dict and k not in pipeline_class._optional_components + } + init_kwargs = {**init_kwargs, **passed_pipe_kwargs} + + # remove `null` components + def load_module(name, value): + if value[0] is None: + return False + if name in passed_class_obj and passed_class_obj[name] is None: + return False + return True + + init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)} + + # Special case: safety_checker must be loaded separately when using `from_flax` + if from_flax and "safety_checker" in init_dict and "safety_checker" not in passed_class_obj: + raise NotImplementedError( + "The safety checker cannot be automatically loaded when loading weights `from_flax`." + " Please, pass `safety_checker=None` to `from_pretrained`, and load the safety checker" + " separately if you need it." + ) + + # 5. Throw nice warnings / errors for fast accelerate loading + if len(unused_kwargs) > 0: + logger.warning( + f"Keyword arguments {unused_kwargs} are not expected by {pipeline_class.__name__} and will be ignored." + ) + + if low_cpu_mem_usage and not is_accelerate_available(): + low_cpu_mem_usage = False + logger.warning( + "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" + " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" + " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" + " install accelerate\n```\n." + ) + + if device_map is not None and not is_torch_version(">=", "1.9.0"): + raise NotImplementedError( + "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set" + " `device_map=None`." + ) + + if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): + raise NotImplementedError( + "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" + " `low_cpu_mem_usage=False`." + ) + + if low_cpu_mem_usage is False and device_map is not None: + raise ValueError( + f"You cannot set `low_cpu_mem_usage` to False while using device_map={device_map} for loading and" + " dispatching. Please make sure to set `low_cpu_mem_usage=True`." + ) + + # import it here to avoid circular import + from diffusers import pipelines + + # 6. Load each module in the pipeline + for name, (library_name, class_name) in tqdm(init_dict.items(), desc="Loading pipeline components..."): + # 6.1 - now that JAX/Flax is an official framework of the library, we might load from Flax names + if class_name.startswith("Flax"): + class_name = class_name[4:] + + # 6.2 Define all importable classes + is_pipeline_module = hasattr(pipelines, library_name) + importable_classes = ALL_IMPORTABLE_CLASSES + loaded_sub_model = None + + # 6.3 Use passed sub model or load class_name from library_name + if name in passed_class_obj: + # if the model is in a pipeline module, then we load it from the pipeline + # check that passed_class_obj has correct parent class + maybe_raise_or_warn( + library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module + ) + + loaded_sub_model = passed_class_obj[name] + else: + # load sub model + loaded_sub_model = load_sub_model( + library_name=library_name, + class_name=class_name, + importable_classes=importable_classes, + pipelines=pipelines, + is_pipeline_module=is_pipeline_module, + pipeline_class=pipeline_class, + torch_dtype=torch_dtype, + provider=provider, + sess_options=sess_options, + device_map=device_map, + max_memory=max_memory, + offload_folder=offload_folder, + offload_state_dict=offload_state_dict, + model_variants=model_variants, + name=name, + from_flax=from_flax, + variant=variant, + low_cpu_mem_usage=low_cpu_mem_usage, + cached_folder=cached_folder, + ) + #logger.info( + # f"Loaded {name} as {class_name} from `{name}` subfolder of {pretrained_model_name_or_path}." + #) + + init_kwargs[name] = loaded_sub_model # UNet(...), # DiffusionSchedule(...) + + if pipeline_class._load_connected_pipes and os.path.isfile(os.path.join(cached_folder, "README.md")): + modelcard = ModelCard.load(os.path.join(cached_folder, "README.md")) + connected_pipes = {prefix: getattr(modelcard.data, prefix, [None])[0] for prefix in CONNECTED_PIPES_KEYS} + load_kwargs = { + "cache_dir": cache_dir, + "resume_download": resume_download, + "force_download": force_download, + "proxies": proxies, + "local_files_only": local_files_only, + "use_auth_token": use_auth_token, + "revision": revision, + "torch_dtype": torch_dtype, + "custom_pipeline": custom_pipeline, + "custom_revision": custom_revision, + "provider": provider, + "sess_options": sess_options, + "device_map": device_map, + "max_memory": max_memory, + "offload_folder": offload_folder, + "offload_state_dict": offload_state_dict, + "low_cpu_mem_usage": low_cpu_mem_usage, + "variant": variant, + "use_safetensors": use_safetensors, + } + + def get_connected_passed_kwargs(prefix): + connected_passed_class_obj = { + k.replace(f"{prefix}_", ""): w for k, w in passed_class_obj.items() if k.split("_")[0] == prefix + } + connected_passed_pipe_kwargs = { + k.replace(f"{prefix}_", ""): w for k, w in passed_pipe_kwargs.items() if k.split("_")[0] == prefix + } + + connected_passed_kwargs = {**connected_passed_class_obj, **connected_passed_pipe_kwargs} + return connected_passed_kwargs + + connected_pipes = { + prefix: DiffusionPipeline.from_pretrained( + repo_id, **load_kwargs.copy(), **get_connected_passed_kwargs(prefix) + ) + for prefix, repo_id in connected_pipes.items() + if repo_id is not None + } + + for prefix, connected_pipe in connected_pipes.items(): + # add connected pipes to `init_kwargs` with _, e.g. "prior_text_encoder" + init_kwargs.update( + {"_".join([prefix, name]): component for name, component in connected_pipe.components.items()} + ) + + # 7. Potentially add passed objects if expected + missing_modules = set(expected_modules) - set(init_kwargs.keys()) + passed_modules = list(passed_class_obj.keys()) + optional_modules = pipeline_class._optional_components + if len(missing_modules) > 0 and missing_modules <= set(passed_modules + optional_modules): + for module in missing_modules: + init_kwargs[module] = passed_class_obj.get(module, None) + elif len(missing_modules) > 0: + passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys())) - optional_kwargs + raise ValueError( + f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed." + ) + + # 8. Instantiate the pipeline + model = pipeline_class(**init_kwargs) + + # 9. Save where the model was instantiated from + model.register_to_config(_name_or_path=pretrained_model_name_or_path) + return model + + @property + def name_or_path(self) -> str: + return getattr(self.config, "_name_or_path", None) + + @property + def _execution_device(self): + r""" + Returns the device on which the pipeline's models will be executed. After calling + [`~DiffusionPipeline.enable_sequential_cpu_offload`] the execution device can only be inferred from + Accelerate's module hooks. + """ + for name, model in self.components.items(): + if not isinstance(model, torch.nn.Module) or name in self._exclude_from_cpu_offload: + continue + + if not hasattr(model, "_hf_hook"): + return self.device + for module in model.modules(): + if ( + hasattr(module, "_hf_hook") + and hasattr(module._hf_hook, "execution_device") + and module._hf_hook.execution_device is not None + ): + return torch.device(module._hf_hook.execution_device) + return self.device + + def enable_model_cpu_offload(self, gpu_id: int = 0, device: Union[torch.device, str] = "cuda"): + r""" + Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared + to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward` + method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with + `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`. + """ + if self.model_cpu_offload_seq is None: + raise ValueError( + "Model CPU offload cannot be enabled because no `model_cpu_offload_seq` class attribute is set." + ) + + if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"): + from accelerate import cpu_offload_with_hook + else: + raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.") + + device = torch.device(f"cuda:{gpu_id}") + + if self.device.type != "cpu": + self.to("cpu", silence_dtype_warnings=True) + device_mod = getattr(torch, self.device.type, None) + if hasattr(device_mod, "empty_cache") and device_mod.is_available(): + device_mod.empty_cache() # otherwise we don't see the memory savings (but they probably exist) + + all_model_components = {k: v for k, v in self.components.items() if isinstance(v, torch.nn.Module)} + + self._all_hooks = [] + hook = None + for model_str in self.model_cpu_offload_seq.split("->"): + model = all_model_components.pop(model_str, None) + if not isinstance(model, torch.nn.Module): + continue + + _, hook = cpu_offload_with_hook(model, device, prev_module_hook=hook) + self._all_hooks.append(hook) + + # CPU offload models that are not in the seq chain unless they are explicitly excluded + # these models will stay on CPU until maybe_free_model_hooks is called + # some models cannot be in the seq chain because they are iteratively called, such as controlnet + for name, model in all_model_components.items(): + if not isinstance(model, torch.nn.Module): + continue + + if name in self._exclude_from_cpu_offload: + model.to(device) + else: + _, hook = cpu_offload_with_hook(model, device) + self._all_hooks.append(hook) + + def maybe_free_model_hooks(self): + r""" + TODO: Better doc string + """ + if not hasattr(self, "_all_hooks") or len(self._all_hooks) == 0: + # `enable_model_cpu_offload` has not be called, so silently do nothing + return + + for hook in self._all_hooks: + # offload model and remove hook from model + hook.offload() + hook.remove() + + # make sure the model is in the same state as before calling it + self.enable_model_cpu_offload() + + def enable_sequential_cpu_offload(self, gpu_id: int = 0, device: Union[torch.device, str] = "cuda"): + r""" + Offloads all models to CPU using 🤗 Accelerate, significantly reducing memory usage. When called, the state + dicts of all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are saved to CPU + and then moved to `torch.device('meta')` and loaded to GPU only when their specific submodule has its `forward` + method called. Offloading happens on a submodule basis. Memory savings are higher than with + `enable_model_cpu_offload`, but performance is lower. + """ + if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"): + from accelerate import cpu_offload + else: + raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher") + + if device == "cuda": + device = torch.device(f"{device}:{gpu_id}") + + if self.device.type != "cpu": + self.to("cpu", silence_dtype_warnings=True) + device_mod = getattr(torch, self.device.type, None) + if hasattr(device_mod, "empty_cache") and device_mod.is_available(): + device_mod.empty_cache() # otherwise we don't see the memory savings (but they probably exist) + + for name, model in self.components.items(): + if not isinstance(model, torch.nn.Module): + continue + + if name in self._exclude_from_cpu_offload: + model.to(device) + else: + # make sure to offload buffers if not all high level weights + # are of type nn.Module + offload_buffers = len(model._parameters) > 0 + cpu_offload(model, device, offload_buffers=offload_buffers) + + @classmethod + def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: + r""" + Download and cache a PyTorch diffusion pipeline from pretrained pipeline weights. + + Parameters: + pretrained_model_name (`str` or `os.PathLike`, *optional*): + A string, the *repository id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline + hosted on the Hub. + custom_pipeline (`str`, *optional*): + Can be either: + + - A string, the *repository id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained + pipeline hosted on the Hub. The repository must contain a file called `pipeline.py` that defines + the custom pipeline. + + - A string, the *file name* of a community pipeline hosted on GitHub under + [Community](https://github.com/huggingface/diffusers/tree/main/examples/community). Valid file + names must match the file name and not the pipeline script (`clip_guided_stable_diffusion` + instead of `clip_guided_stable_diffusion.py`). Community pipelines are always loaded from the + current `main` branch of GitHub. + + - A path to a *directory* (`./my_pipeline_directory/`) containing a custom pipeline. The directory + must contain a file called `pipeline.py` that defines the custom pipeline. + + + + 🧪 This is an experimental feature and may change in the future. + + + + For more information on how to load and create custom pipelines, take a look at [How to contribute a + community pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/contribute_pipeline). + + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + resume_download (`bool`, *optional*, defaults to `False`): + Whether or not to resume downloading the model weights and configuration files. If set to `False`, any + incompletely downloaded files are deleted. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info(`bool`, *optional*, defaults to `False`): + Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files or not. If set to `True`, the model + won't be downloaded from the Hub. + use_auth_token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from + `diffusers-cli login` (stored in `~/.huggingface`) is used. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier + allowed by Git. + custom_revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id similar to + `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a + custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. + mirror (`str`, *optional*): + Mirror source to resolve accessibility issues if you're downloading a model in China. We do not + guarantee the timeliness or safety of the source, and you should refer to the mirror site for more + information. + variant (`str`, *optional*): + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when + loading `from_flax`. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `None`, the safetensors weights are downloaded if they're available **and** if the + safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors + weights. If set to `False`, safetensors weights are not loaded. + use_onnx (`bool`, *optional*, defaults to `False`): + If set to `True`, ONNX weights will always be downloaded if present. If set to `False`, ONNX weights + will never be downloaded. By default `use_onnx` defaults to the `_is_onnx` class attribute which is + `False` for non-ONNX pipelines and `True` for ONNX pipelines. ONNX weights include both files ending + with `.onnx` and `.pb`. + + Returns: + `os.PathLike`: + A path to the downloaded pipeline. + + + + To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with + `huggingface-cli login`. + + + + """ + cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE) + resume_download = kwargs.pop("resume_download", False) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + from_flax = kwargs.pop("from_flax", False) + custom_pipeline = kwargs.pop("custom_pipeline", None) + custom_revision = kwargs.pop("custom_revision", None) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None) + use_onnx = kwargs.pop("use_onnx", None) + load_connected_pipeline = kwargs.pop("load_connected_pipeline", False) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = True + allow_pickle = True + + allow_patterns = None + ignore_patterns = None + + model_info_call_error: Optional[Exception] = None + if not local_files_only: + try: + info = model_info( + pretrained_model_name, + use_auth_token=use_auth_token, + revision=revision, + ) + except HTTPError as e: + logger.warn(f"Couldn't connect to the Hub: {e}.\nWill try to load from local cache.") + local_files_only = True + model_info_call_error = e # save error to reraise it if model is not cached locally + + if not local_files_only: + config_file = hf_hub_download( + pretrained_model_name, + cls.config_name, + cache_dir=cache_dir, + revision=revision, + proxies=proxies, + force_download=force_download, + resume_download=resume_download, + use_auth_token=use_auth_token, + ) + + config_dict = cls._dict_from_json_file(config_file) + + ignore_filenames = config_dict.pop("_ignore_files", []) + + # retrieve all folder_names that contain relevant files + folder_names = [k for k, v in config_dict.items() if isinstance(v, list)] + + filenames = {sibling.rfilename for sibling in info.siblings} + model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant) + + if len(variant_filenames) == 0 and variant is not None: + deprecation_message = ( + f"You are trying to load the model files of the `variant={variant}`, but no such modeling files are available." + f"The default model files: {model_filenames} will be loaded instead. Make sure to not load from `variant={variant}`" + "if such variant modeling files are not available. Doing so will lead to an error in v0.22.0 as defaulting to non-variant" + "modeling files is deprecated." + ) + deprecate("no variant default", "0.22.0", deprecation_message, standard_warn=False) + + # remove ignored filenames + model_filenames = set(model_filenames) - set(ignore_filenames) + variant_filenames = set(variant_filenames) - set(ignore_filenames) + + # if the whole pipeline is cached we don't have to ping the Hub + if revision in DEPRECATED_REVISION_ARGS and version.parse( + version.parse(__version__).base_version + ) >= version.parse("0.22.0"): + warn_deprecated_model_variant( + pretrained_model_name, use_auth_token, variant, revision, model_filenames + ) + + model_folder_names = {os.path.split(f)[0] for f in model_filenames if os.path.split(f)[0] in folder_names} + + # all filenames compatible with variant will be added + allow_patterns = list(model_filenames) + + # allow all patterns from non-model folders + # this enables downloading schedulers, tokenizers, ... + allow_patterns += [f"{k}/*" for k in folder_names if k not in model_folder_names] + # also allow downloading config.json files with the model + allow_patterns += [os.path.join(k, "config.json") for k in model_folder_names] + + allow_patterns += [ + SCHEDULER_CONFIG_NAME, + CONFIG_NAME, + cls.config_name, + CUSTOM_PIPELINE_FILE_NAME, + ] + + # retrieve passed components that should not be downloaded + pipeline_class = _get_pipeline_class( + cls, + config_dict, + load_connected_pipeline=load_connected_pipeline, + custom_pipeline=custom_pipeline, + cache_dir=cache_dir, + revision=custom_revision, + ) + expected_components, _ = cls._get_signature_keys(pipeline_class) + passed_components = [k for k in expected_components if k in kwargs] + + if ( + use_safetensors + and not allow_pickle + and not is_safetensors_compatible( + model_filenames, variant=variant, passed_components=passed_components + ) + ): + raise EnvironmentError( + f"Could not found the necessary `safetensors` weights in {model_filenames} (variant={variant})" + ) + if from_flax: + ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb"] + elif use_safetensors and is_safetensors_compatible( + model_filenames, variant=variant, passed_components=passed_components + ): + ignore_patterns = ["*.bin", "*.msgpack"] + + use_onnx = use_onnx if use_onnx is not None else pipeline_class._is_onnx + if not use_onnx: + ignore_patterns += ["*.onnx", "*.pb"] + + safetensors_variant_filenames = {f for f in variant_filenames if f.endswith(".safetensors")} + safetensors_model_filenames = {f for f in model_filenames if f.endswith(".safetensors")} + if ( + len(safetensors_variant_filenames) > 0 + and safetensors_model_filenames != safetensors_variant_filenames + ): + logger.warn( + f"\nA mixture of {variant} and non-{variant} filenames will be loaded.\nLoaded {variant} filenames:\n[{', '.join(safetensors_variant_filenames)}]\nLoaded non-{variant} filenames:\n[{', '.join(safetensors_model_filenames - safetensors_variant_filenames)}\nIf this behavior is not expected, please check your folder structure." + ) + else: + ignore_patterns = ["*.safetensors", "*.msgpack"] + + use_onnx = use_onnx if use_onnx is not None else pipeline_class._is_onnx + if not use_onnx: + ignore_patterns += ["*.onnx", "*.pb"] + + bin_variant_filenames = {f for f in variant_filenames if f.endswith(".bin")} + bin_model_filenames = {f for f in model_filenames if f.endswith(".bin")} + if len(bin_variant_filenames) > 0 and bin_model_filenames != bin_variant_filenames: + logger.warn( + f"\nA mixture of {variant} and non-{variant} filenames will be loaded.\nLoaded {variant} filenames:\n[{', '.join(bin_variant_filenames)}]\nLoaded non-{variant} filenames:\n[{', '.join(bin_model_filenames - bin_variant_filenames)}\nIf this behavior is not expected, please check your folder structure." + ) + + # Don't download any objects that are passed + allow_patterns = [ + p for p in allow_patterns if not (len(p.split("/")) == 2 and p.split("/")[0] in passed_components) + ] + + if pipeline_class._load_connected_pipes: + allow_patterns.append("README.md") + + # Don't download index files of forbidden patterns either + ignore_patterns = ignore_patterns + [f"{i}.index.*json" for i in ignore_patterns] + + re_ignore_pattern = [re.compile(fnmatch.translate(p)) for p in ignore_patterns] + re_allow_pattern = [re.compile(fnmatch.translate(p)) for p in allow_patterns] + + expected_files = [f for f in filenames if not any(p.match(f) for p in re_ignore_pattern)] + expected_files = [f for f in expected_files if any(p.match(f) for p in re_allow_pattern)] + + snapshot_folder = Path(config_file).parent + pipeline_is_cached = all((snapshot_folder / f).is_file() for f in expected_files) + + if pipeline_is_cached and not force_download: + # if the pipeline is cached, we can directly return it + # else call snapshot_download + return snapshot_folder + + user_agent = {"pipeline_class": cls.__name__} + if custom_pipeline is not None and not custom_pipeline.endswith(".py"): + user_agent["custom_pipeline"] = custom_pipeline + + # download all allow_patterns - ignore_patterns + try: + cached_folder = snapshot_download( + pretrained_model_name, + cache_dir=cache_dir, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + user_agent=user_agent, + ) + + # retrieve pipeline class from local file + cls_name = cls.load_config(os.path.join(cached_folder, "model_index.json")).get("_class_name", None) + pipeline_class = getattr(diffusers, cls_name, None) + + if pipeline_class is not None and pipeline_class._load_connected_pipes: + modelcard = ModelCard.load(os.path.join(cached_folder, "README.md")) + connected_pipes = sum([getattr(modelcard.data, k, []) for k in CONNECTED_PIPES_KEYS], []) + for connected_pipe_repo_id in connected_pipes: + download_kwargs = { + "cache_dir": cache_dir, + "resume_download": resume_download, + "force_download": force_download, + "proxies": proxies, + "local_files_only": local_files_only, + "use_auth_token": use_auth_token, + "variant": variant, + "use_safetensors": use_safetensors, + } + DiffusionPipeline.download(connected_pipe_repo_id, **download_kwargs) + + return cached_folder + + except FileNotFoundError: + # Means we tried to load pipeline with `local_files_only=True` but the files have not been found in local cache. + # This can happen in two cases: + # 1. If the user passed `local_files_only=True` => we raise the error directly + # 2. If we forced `local_files_only=True` when `model_info` failed => we raise the initial error + if model_info_call_error is None: + # 1. user passed `local_files_only=True` + raise + else: + # 2. we forced `local_files_only=True` when `model_info` failed + raise EnvironmentError( + f"Cannot load model {pretrained_model_name}: model is not cached locally and an error occured" + " while trying to fetch metadata from the Hub. Please check out the root cause in the stacktrace" + " above." + ) from model_info_call_error + + @staticmethod + def _get_signature_keys(obj): + parameters = inspect.signature(obj.__init__).parameters + required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty} + optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty}) + expected_modules = set(required_parameters.keys()) - {"self"} + return expected_modules, optional_parameters + + @property + def components(self) -> Dict[str, Any]: + r""" + The `self.components` property can be useful to run different pipelines with the same weights and + configurations without reallocating additional memory. + + Returns (`dict`): + A dictionary containing all the modules needed to initialize the pipeline. + + Examples: + + ```py + >>> from diffusers import ( + ... StableDiffusionPipeline, + ... StableDiffusionImg2ImgPipeline, + ... StableDiffusionInpaintPipeline, + ... ) + + >>> text2img = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> img2img = StableDiffusionImg2ImgPipeline(**text2img.components) + >>> inpaint = StableDiffusionInpaintPipeline(**text2img.components) + ``` + """ + expected_modules, optional_parameters = self._get_signature_keys(self) + components = { + k: getattr(self, k) for k in self.config.keys() if not k.startswith("_") and k not in optional_parameters + } + + if set(components.keys()) != expected_modules: + raise ValueError( + f"{self} has been incorrectly initialized or {self.__class__} is incorrectly implemented. Expected" + f" {expected_modules} to be defined, but {components.keys()} are defined." + ) + + return components + + @staticmethod + def numpy_to_pil(images): + """ + Convert a NumPy image or a batch of images to a PIL image. + """ + return numpy_to_pil(images) + + def progress_bar(self, iterable=None, total=None): + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + elif not isinstance(self._progress_bar_config, dict): + raise ValueError( + f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." + ) + + if iterable is not None: + return tqdm(iterable, **self._progress_bar_config) + elif total is not None: + return tqdm(total=total, **self._progress_bar_config) + else: + raise ValueError("Either `total` or `iterable` has to be defined.") + + def set_progress_bar_config(self, **kwargs): + self._progress_bar_config = kwargs + + def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + r""" + Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). When this + option is enabled, you should observe lower GPU memory usage and a potential speed up during inference. Speed + up during training is not guaranteed. + + + + ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes + precedent. + + + + Parameters: + attention_op (`Callable`, *optional*): + Override the default `None` operator for use as `op` argument to the + [`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention) + function of xFormers. + + Examples: + + ```py + >>> import torch + >>> from diffusers import DiffusionPipeline + >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp + + >>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16) + >>> pipe = pipe.to("cuda") + >>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp) + >>> # Workaround for not accepting attention shape using VAE for Flash Attention + >>> pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None) + ``` + """ + self.set_use_memory_efficient_attention_xformers(True, attention_op) + + def disable_xformers_memory_efficient_attention(self): + r""" + Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). + """ + self.set_use_memory_efficient_attention_xformers(False) + + def set_use_memory_efficient_attention_xformers( + self, valid: bool, attention_op: Optional[Callable] = None + ) -> None: + # Recursively walk through all the children. + # Any children which exposes the set_use_memory_efficient_attention_xformers method + # gets the message + def fn_recursive_set_mem_eff(module: torch.nn.Module): + if hasattr(module, "set_use_memory_efficient_attention_xformers"): + module.set_use_memory_efficient_attention_xformers(valid, attention_op) + + for child in module.children(): + fn_recursive_set_mem_eff(child) + + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + for module in modules: + fn_recursive_set_mem_eff(module) + + def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"): + r""" + Enable sliced attention computation. When this option is enabled, the attention module splits the input tensor + in slices to compute attention in several steps. For more than one attention head, the computation is performed + sequentially over each head. This is useful to save some memory in exchange for a small speed decrease. + + + + ⚠️ Don't enable attention slicing if you're already using `scaled_dot_product_attention` (SDPA) from PyTorch + 2.0 or xFormers. These attention computations are already very memory efficient so you won't need to enable + this function. If you enable attention slicing with SDPA or xFormers, it can lead to serious slow downs! + + + + Args: + slice_size (`str` or `int`, *optional*, defaults to `"auto"`): + When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If + `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + + Examples: + + ```py + >>> import torch + >>> from diffusers import StableDiffusionPipeline + + >>> pipe = StableDiffusionPipeline.from_pretrained( + ... "runwayml/stable-diffusion-v1-5", + ... torch_dtype=torch.float16, + ... use_safetensors=True, + ... ) + + >>> prompt = "a photo of an astronaut riding a horse on mars" + >>> pipe.enable_attention_slicing() + >>> image = pipe(prompt).images[0] + ``` + """ + self.set_attention_slice(slice_size) + + def disable_attention_slicing(self): + r""" + Disable sliced attention computation. If `enable_attention_slicing` was previously called, attention is + computed in one step. + """ + # set slice_size = `None` to disable `attention slicing` + self.enable_attention_slicing(None) + + def set_attention_slice(self, slice_size: Optional[int]): + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module) and hasattr(m, "set_attention_slice")] + + for module in modules: + module.set_attention_slice(slice_size) diff --git a/ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_blocks.py b/ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_blocks.py new file mode 100644 index 00000000..81036c63 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_blocks.py @@ -0,0 +1,3339 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. +from typing import Any, Dict, Optional, Tuple + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from diffusers.utils import is_torch_version, logging +from diffusers.models.activations import get_activation +import diffusers +if diffusers.__version__ >= '0.22.0': + from diffusers.models.normalization import AdaGroupNorm +else: + from diffusers.models.attention import AdaGroupNorm +from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0 +from diffusers.models.dual_transformer_2d import DualTransformer2DModel +from diffusers.models.resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, Upsample2D +from diffusers.models.transformer_2d import Transformer2DModel + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +import time + +def get_down_block( + down_block_type, + num_layers, + in_channels, + out_channels, + temb_channels, + add_downsample, + resnet_eps, + resnet_act_fn, + transformer_layers_per_block=1, + num_attention_heads=None, + resnet_groups=None, + cross_attention_dim=None, + downsample_padding=None, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + resnet_time_scale_shift="default", + attention_type="default", + resnet_skip_time_act=False, + resnet_out_scale_factor=1.0, + cross_attention_norm=None, + attention_head_dim=None, + downsample_type=None, + dropout=0.0, +): + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type + if down_block_type == "DownBlock2D": + return DownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "ResnetDownsampleBlock2D": + return ResnetDownsampleBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + ) + elif down_block_type == "AttnDownBlock2D": + if add_downsample is False: + downsample_type = None + else: + downsample_type = downsample_type or "conv" # default to 'conv' + return AttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + downsample_type=downsample_type, + ) + elif down_block_type == "CrossAttnDownBlock2D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D") + return CrossAttnDownBlock2D( + num_layers=num_layers, + transformer_layers_per_block=transformer_layers_per_block, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + ) + elif down_block_type == "SimpleCrossAttnDownBlock2D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D") + return SimpleCrossAttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif down_block_type == "SkipDownBlock2D": + return SkipDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "AttnSkipDownBlock2D": + return AttnSkipDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "DownEncoderBlock2D": + return DownEncoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "AttnDownEncoderBlock2D": + return AttnDownEncoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "KDownBlock2D": + return KDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + ) + elif down_block_type == "KCrossAttnDownBlock2D": + return KCrossAttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + cross_attention_dim=cross_attention_dim, + attention_head_dim=attention_head_dim, + add_self_attention=True if not add_downsample else False, + ) + raise ValueError(f"{down_block_type} does not exist.") + + +def get_up_block( + up_block_type, + num_layers, + in_channels, + out_channels, + prev_output_channel, + temb_channels, + add_upsample, + resnet_eps, + resnet_act_fn, + transformer_layers_per_block=1, + num_attention_heads=None, + resnet_groups=None, + cross_attention_dim=None, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + resnet_time_scale_shift="default", + attention_type="default", + resnet_skip_time_act=False, + resnet_out_scale_factor=1.0, + cross_attention_norm=None, + attention_head_dim=None, + upsample_type=None, + dropout=0.0, +): + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type + if up_block_type == "UpBlock2D": + return UpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == "ResnetUpsampleBlock2D": + return ResnetUpsampleBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + ) + elif up_block_type == "CrossAttnUpBlock2D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D") + return CrossAttnUpBlock2D( + num_layers=num_layers, + transformer_layers_per_block=transformer_layers_per_block, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + ) + elif up_block_type == "SimpleCrossAttnUpBlock2D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D") + return SimpleCrossAttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif up_block_type == "AttnUpBlock2D": + if add_upsample is False: + upsample_type = None + else: + upsample_type = upsample_type or "conv" # default to 'conv' + + return AttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + upsample_type=upsample_type, + ) + elif up_block_type == "SkipUpBlock2D": + return SkipUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == "AttnSkipUpBlock2D": + return AttnSkipUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == "UpDecoderBlock2D": + return UpDecoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + temb_channels=temb_channels, + ) + elif up_block_type == "AttnUpDecoderBlock2D": + return AttnUpDecoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + attention_head_dim=attention_head_dim, + resnet_time_scale_shift=resnet_time_scale_shift, + temb_channels=temb_channels, + ) + elif up_block_type == "KUpBlock2D": + return KUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + ) + elif up_block_type == "KCrossAttnUpBlock2D": + return KCrossAttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + cross_attention_dim=cross_attention_dim, + attention_head_dim=attention_head_dim, + ) + + raise ValueError(f"{up_block_type} does not exist.") + + +class AutoencoderTinyBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int, act_fn: str): + super().__init__() + act_fn = get_activation(act_fn) + self.conv = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), + act_fn, + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + act_fn, + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + ) + self.skip = ( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) + if in_channels != out_channels + else nn.Identity() + ) + self.fuse = nn.ReLU() + + def forward(self, x): + return self.fuse(self.conv(x) + self.skip(x)) + + +class UNetMidBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + ): + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + self.add_attention = add_attention + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." + ) + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups if resnet_time_scale_shift == "default" else None, + spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + else: + attentions.append(None) + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, hidden_states, temb=None): + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + hidden_states = attn(hidden_states, temb=temb) + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class UNetMidBlock2DCrossAttn(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads=1, + output_scale_factor=1.0, + cross_attention_dim=1280, + dual_cross_attention=False, + use_linear_projection=False, + upcast_attention=False, + attention_type="default", + ): + super().__init__() + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + for _ in range(num_layers): + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + in_channels // num_attention_heads, + in_channels=in_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + in_channels // num_attention_heads, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + else: + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + return hidden_states + + +class UNetMidBlock2DSimpleCrossAttn(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + cross_attention_dim=1280, + skip_time_act=False, + only_cross_attention=False, + cross_attention_norm=None, + ): + super().__init__() + + self.has_cross_attention = True + + self.attention_head_dim = attention_head_dim + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + + self.num_heads = in_channels // self.attention_head_dim + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ] + attentions = [] + + for _ in range(num_layers): + processor = ( + AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() + ) + + attentions.append( + Attention( + query_dim=in_channels, + cross_attention_dim=in_channels, + heads=self.num_heads, + dim_head=self.attention_head_dim, + added_kv_proj_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + bias=True, + upcast_softmax=True, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + processor=processor, + ) + ) + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + lora_scale = cross_attention_kwargs.get("scale", 1.0) + + if attention_mask is None: + # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. + mask = None if encoder_hidden_states is None else encoder_attention_mask + else: + # when attention_mask is defined: we don't even check for encoder_attention_mask. + # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. + # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. + # then we can simplify this whole if/else block to: + # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask + mask = attention_mask + + hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + # attn + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + + # resnet + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + return hidden_states + + +class AttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + downsample_padding=1, + downsample_type="conv", + ): + super().__init__() + resnets = [] + attentions = [] + self.downsample_type = downsample_type + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if downsample_type == "conv": + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + elif downsample_type == "resnet": + self.downsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + down=True, + ) + ] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states, temb=None, upsample_size=None, cross_attention_kwargs=None): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + lora_scale = cross_attention_kwargs.get("scale", 1.0) + + output_states = () + + for resnet, attn in zip(self.resnets, self.attentions): + cross_attention_kwargs.update({"scale": lora_scale}) + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn(hidden_states, **cross_attention_kwargs) + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + if self.downsample_type == "resnet": + hidden_states = downsampler(hidden_states, temb=temb, scale=lora_scale) + else: + hidden_states = downsampler(hidden_states, scale=lora_scale) + + output_states += (hidden_states,) + + return hidden_states, output_states + + +class CrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + downsample_padding=1, + add_downsample=True, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + attention_type="default", + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + exist_block_number=None, + additional_residuals=None, + ): + output_states = () + + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + blocks = list(zip(self.resnets, self.attentions)) + + for i, (resnet, attn) in enumerate(blocks): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + # apply additional residuals to the output of the last pair of resnet and attention blocks + if i == len(blocks) - 1 and additional_residuals is not None: + hidden_states = hidden_states + additional_residuals + + output_states = output_states + (hidden_states,) + if exist_block_number is not None and len(output_states) == exist_block_number + 1: + return hidden_states, output_states + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale=lora_scale) + + output_states = output_states + (hidden_states,) + return hidden_states, output_states + + +class DownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None, scale: float = 1.0, exist_block_number=None,): + output_states = () + + i = 0 + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + output_states = output_states + (hidden_states,) + if exist_block_number is not None and len(output_states) == exist_block_number + 1: + return hidden_states, output_states + i += 1 + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale=scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class DownEncoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states, scale: float = 1.0): + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None, scale=scale) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale) + + return hidden_states + + +class AttnDownEncoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" + ) + ] + ) + else: + self.downsamplers = None + + def forward(self, hidden_states, scale: float = 1.0): + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb=None, scale=scale) + cross_attention_kwargs = {"scale": scale} + hidden_states = attn(hidden_states, **cross_attention_kwargs) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale) + + return hidden_states + + +class AttnSkipDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=np.sqrt(2.0), + add_downsample=True, + ): + super().__init__() + self.attentions = nn.ModuleList([]) + self.resnets = nn.ModuleList([]) + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(in_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + self.attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=32, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + if add_downsample: + self.resnet_down = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + down=True, + kernel="fir", + ) + self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)]) + self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1)) + else: + self.resnet_down = None + self.downsamplers = None + self.skip_conv = None + + def forward(self, hidden_states, temb=None, skip_sample=None, scale: float = 1.0): + output_states = () + + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb, scale=scale) + cross_attention_kwargs = {"scale": scale} + hidden_states = attn(hidden_states, **cross_attention_kwargs) + output_states += (hidden_states,) + + if self.downsamplers is not None: + hidden_states = self.resnet_down(hidden_states, temb, scale=scale) + for downsampler in self.downsamplers: + skip_sample = downsampler(skip_sample) + + hidden_states = self.skip_conv(skip_sample) + hidden_states + + output_states += (hidden_states,) + + return hidden_states, output_states, skip_sample + + +class SkipDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_pre_norm: bool = True, + output_scale_factor=np.sqrt(2.0), + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(in_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + if add_downsample: + self.resnet_down = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + down=True, + kernel="fir", + ) + self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)]) + self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1)) + else: + self.resnet_down = None + self.downsamplers = None + self.skip_conv = None + + def forward(self, hidden_states, temb=None, skip_sample=None, scale: float = 1.0): + output_states = () + + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb, scale) + output_states += (hidden_states,) + + if self.downsamplers is not None: + hidden_states = self.resnet_down(hidden_states, temb, scale) + for downsampler in self.downsamplers: + skip_sample = downsampler(skip_sample) + + hidden_states = self.skip_conv(skip_sample) + hidden_states + + output_states += (hidden_states,) + + return hidden_states, output_states, skip_sample + + +class ResnetDownsampleBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + skip_time_act=False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + down=True, + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None, scale: float = 1.0): + output_states = () + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale) + + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, temb, scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class SimpleCrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + add_downsample=True, + skip_time_act=False, + only_cross_attention=False, + cross_attention_norm=None, + ): + super().__init__() + + self.has_cross_attention = True + + resnets = [] + attentions = [] + + self.attention_head_dim = attention_head_dim + self.num_heads = out_channels // self.attention_head_dim + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + processor = ( + AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() + ) + + attentions.append( + Attention( + query_dim=out_channels, + cross_attention_dim=out_channels, + heads=self.num_heads, + dim_head=attention_head_dim, + added_kv_proj_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + bias=True, + upcast_softmax=True, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + processor=processor, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + down=True, + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + output_states = () + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + lora_scale = cross_attention_kwargs.get("scale", 1.0) + + if attention_mask is None: + # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. + mask = None if encoder_hidden_states is None else encoder_attention_mask + else: + # when attention_mask is defined: we don't even check for encoder_attention_mask. + # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. + # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. + # then we can simplify this whole if/else block to: + # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask + mask = attention_mask + + for resnet, attn in zip(self.resnets, self.attentions): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, temb, scale=lora_scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class KDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 4, + resnet_eps: float = 1e-5, + resnet_act_fn: str = "gelu", + resnet_group_size: int = 32, + add_downsample=False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + temb_channels=temb_channels, + groups=groups, + groups_out=groups_out, + eps=resnet_eps, + non_linearity=resnet_act_fn, + time_embedding_norm="ada_group", + conv_shortcut_bias=False, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + # YiYi's comments- might be able to use FirDownsample2D, look into details later + self.downsamplers = nn.ModuleList([KDownsample2D()]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None, scale: float = 1.0): + output_states = () + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale) + + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states, output_states + + +class KCrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + cross_attention_dim: int, + dropout: float = 0.0, + num_layers: int = 4, + resnet_group_size: int = 32, + add_downsample=True, + attention_head_dim: int = 64, + add_self_attention: bool = False, + resnet_eps: float = 1e-5, + resnet_act_fn: str = "gelu", + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + temb_channels=temb_channels, + groups=groups, + groups_out=groups_out, + eps=resnet_eps, + non_linearity=resnet_act_fn, + time_embedding_norm="ada_group", + conv_shortcut_bias=False, + ) + ) + attentions.append( + KAttentionBlock( + out_channels, + out_channels // attention_head_dim, + attention_head_dim, + cross_attention_dim=cross_attention_dim, + temb_channels=temb_channels, + attention_bias=True, + add_self_attention=add_self_attention, + cross_attention_norm="layer_norm", + group_size=resnet_group_size, + ) + ) + + self.resnets = nn.ModuleList(resnets) + self.attentions = nn.ModuleList(attentions) + + if add_downsample: + self.downsamplers = nn.ModuleList([KDownsample2D()]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + output_states = () + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + for resnet, attn in zip(self.resnets, self.attentions): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + + if self.downsamplers is None: + output_states += (None,) + else: + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states, output_states + + +class AttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + upsample_type="conv", + ): + super().__init__() + resnets = [] + attentions = [] + + self.upsample_type = upsample_type + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if upsample_type == "conv": + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + elif upsample_type == "resnet": + self.upsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + up=True, + ) + ] + ) + else: + self.upsamplers = None + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + for resnet, attn in zip(self.resnets, self.attentions): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb, scale=scale) + cross_attention_kwargs = {"scale": scale} + hidden_states = attn(hidden_states, **cross_attention_kwargs) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + if self.upsample_type == "resnet": + hidden_states = upsampler(hidden_states, temb=temb, scale=scale) + else: + hidden_states = upsampler(hidden_states, scale=scale) + + return hidden_states + + +class CrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + add_upsample=True, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + attention_type="default", + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + enter_block_number: Optional[int]=None, + ): + prv_f = [] + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)): + # pop res hidden states + + if enter_block_number is not None and i < len(self.resnets) - enter_block_number - 1: + continue + + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + prv_f.append(hidden_states) + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + if not hidden_states.requires_grad and hidden_states.shape[0]>=64: + _chunk_size=8 + num_chunks = hidden_states.shape[0] // _chunk_size + hidden_states = torch.cat( + [ + upsampler(hid_slice, upsample_size, scale=lora_scale) + + for hid_slice in hidden_states.chunk(num_chunks, dim=0) + ], + dim=0, + ) + else: + hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale) + + return hidden_states, prv_f + + +class UpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0, enter_block_number: Optional[int]=None,): + prv_f = [] + + for idx, resnet in enumerate(self.resnets): + + if enter_block_number is not None and idx < len(self.resnets) - enter_block_number - 1: + continue + + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + prv_f.append(hidden_states) + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + if not hidden_states.requires_grad and hidden_states.shape[0]>=64: + _chunk_size=8 + num_chunks = hidden_states.shape[0] // _chunk_size + hidden_states = torch.cat( + [ + resnet(hid_slice, temb_slice, scale=scale) + + for hid_slice,temb_slice in zip(hidden_states.chunk(num_chunks, dim=0),temb.chunk(num_chunks, dim=0)) + ], + dim=0, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + if not hidden_states.requires_grad and hidden_states.shape[0]>=64: + _chunk_size=8 + num_chunks = hidden_states.shape[0] // _chunk_size + hidden_states = torch.cat( + [ + upsampler(hid_slice, upsample_size, scale=scale) + + for hid_slice in hidden_states.chunk(num_chunks, dim=0) + ], + dim=0, + ) + else: + hidden_states = upsampler(hidden_states, upsample_size, scale=scale) + + return hidden_states, prv_f + + +class UpDecoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + temb_channels=None, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + def forward(self, hidden_states, temb=None, scale: float = 1.0): + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class AttnUpDecoderBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=1.0, + add_upsample=True, + temb_channels=None, + ): + super().__init__() + resnets = [] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups if resnet_time_scale_shift != "spatial" else None, + spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + def forward(self, hidden_states, temb=None, scale: float = 1.0): + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb=temb, scale=scale) + cross_attention_kwargs = {"scale": scale} + hidden_states = attn(hidden_states, temb=temb, **cross_attention_kwargs) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, scale=scale) + + return hidden_states + + +class AttnSkipUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_pre_norm: bool = True, + attention_head_dim=1, + output_scale_factor=np.sqrt(2.0), + add_upsample=True, + ): + super().__init__() + self.attentions = nn.ModuleList([]) + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + self.resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(resnet_in_channels + res_skip_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}." + ) + attention_head_dim = out_channels + + self.attentions.append( + Attention( + out_channels, + heads=out_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=32, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + + self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) + if add_upsample: + self.resnet_up = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + up=True, + kernel="fir", + ) + self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + self.skip_norm = torch.nn.GroupNorm( + num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True + ) + self.act = nn.SiLU() + else: + self.resnet_up = None + self.skip_conv = None + self.skip_norm = None + self.act = None + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=None, scale: float = 1.0): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb, scale=scale) + + cross_attention_kwargs = {"scale": scale} + hidden_states = self.attentions[0](hidden_states, **cross_attention_kwargs) + + if skip_sample is not None: + skip_sample = self.upsampler(skip_sample) + else: + skip_sample = 0 + + if self.resnet_up is not None: + skip_sample_states = self.skip_norm(hidden_states) + skip_sample_states = self.act(skip_sample_states) + skip_sample_states = self.skip_conv(skip_sample_states) + + skip_sample = skip_sample + skip_sample_states + + hidden_states = self.resnet_up(hidden_states, temb, scale=scale) + + return hidden_states, skip_sample + + +class SkipUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_pre_norm: bool = True, + output_scale_factor=np.sqrt(2.0), + add_upsample=True, + upsample_padding=1, + ): + super().__init__() + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + self.resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min((resnet_in_channels + res_skip_channels) // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) + if add_upsample: + self.resnet_up = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + up=True, + kernel="fir", + ) + self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) + self.skip_norm = torch.nn.GroupNorm( + num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True + ) + self.act = nn.SiLU() + else: + self.resnet_up = None + self.skip_conv = None + self.skip_norm = None + self.act = None + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=None, scale: float = 1.0): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb, scale=scale) + + if skip_sample is not None: + skip_sample = self.upsampler(skip_sample) + else: + skip_sample = 0 + + if self.resnet_up is not None: + skip_sample_states = self.skip_norm(hidden_states) + skip_sample_states = self.act(skip_sample_states) + skip_sample_states = self.skip_conv(skip_sample_states) + + skip_sample = skip_sample + skip_sample_states + + hidden_states = self.resnet_up(hidden_states, temb, scale=scale) + + return hidden_states, skip_sample + + +class ResnetUpsampleBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + skip_time_act=False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + up=True, + ) + ] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, temb, scale=scale) + + return hidden_states + + +class SimpleCrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attention_head_dim=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + add_upsample=True, + skip_time_act=False, + only_cross_attention=False, + cross_attention_norm=None, + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.attention_head_dim = attention_head_dim + + self.num_heads = out_channels // self.attention_head_dim + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ) + + processor = ( + AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor() + ) + + attentions.append( + Attention( + query_dim=out_channels, + cross_attention_dim=out_channels, + heads=self.num_heads, + dim_head=self.attention_head_dim, + added_kv_proj_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + bias=True, + upcast_softmax=True, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + processor=processor, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + up=True, + ) + ] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + lora_scale = cross_attention_kwargs.get("scale", 1.0) + if attention_mask is None: + # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask. + mask = None if encoder_hidden_states is None else encoder_attention_mask + else: + # when attention_mask is defined: we don't even check for encoder_attention_mask. + # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks. + # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask. + # then we can simplify this whole if/else block to: + # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask + mask = attention_mask + + for resnet, attn in zip(self.resnets, self.attentions): + # resnet + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=mask, + **cross_attention_kwargs, + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, temb, scale=lora_scale) + + return hidden_states + + +class KUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 5, + resnet_eps: float = 1e-5, + resnet_act_fn: str = "gelu", + resnet_group_size: Optional[int] = 32, + add_upsample=True, + ): + super().__init__() + resnets = [] + k_in_channels = 2 * out_channels + k_out_channels = in_channels + num_layers = num_layers - 1 + + for i in range(num_layers): + in_channels = k_in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=k_out_channels if (i == num_layers - 1) else out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=groups, + groups_out=groups_out, + dropout=dropout, + non_linearity=resnet_act_fn, + time_embedding_norm="ada_group", + conv_shortcut_bias=False, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([KUpsample2D()]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, scale: float = 1.0): + res_hidden_states_tuple = res_hidden_states_tuple[-1] + if res_hidden_states_tuple is not None: + hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1) + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, use_reentrant=False + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class KCrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 4, + resnet_eps: float = 1e-5, + resnet_act_fn: str = "gelu", + resnet_group_size: int = 32, + attention_head_dim=1, # attention dim_head + cross_attention_dim: int = 768, + add_upsample: bool = True, + upcast_attention: bool = False, + ): + super().__init__() + resnets = [] + attentions = [] + + is_first_block = in_channels == out_channels == temb_channels + is_middle_block = in_channels != out_channels + add_self_attention = True if is_first_block else False + + self.has_cross_attention = True + self.attention_head_dim = attention_head_dim + + # in_channels, and out_channels for the block (k-unet) + k_in_channels = out_channels if is_first_block else 2 * out_channels + k_out_channels = in_channels + + num_layers = num_layers - 1 + + for i in range(num_layers): + in_channels = k_in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + if is_middle_block and (i == num_layers - 1): + conv_2d_out_channels = k_out_channels + else: + conv_2d_out_channels = None + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + conv_2d_out_channels=conv_2d_out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=groups, + groups_out=groups_out, + dropout=dropout, + non_linearity=resnet_act_fn, + time_embedding_norm="ada_group", + conv_shortcut_bias=False, + ) + ) + attentions.append( + KAttentionBlock( + k_out_channels if (i == num_layers - 1) else out_channels, + k_out_channels // attention_head_dim + if (i == num_layers - 1) + else out_channels // attention_head_dim, + attention_head_dim, + cross_attention_dim=cross_attention_dim, + temb_channels=temb_channels, + attention_bias=True, + add_self_attention=add_self_attention, + cross_attention_norm="layer_norm", + upcast_attention=upcast_attention, + ) + ) + + self.resnets = nn.ModuleList(resnets) + self.attentions = nn.ModuleList(attentions) + + if add_upsample: + self.upsamplers = nn.ModuleList([KUpsample2D()]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + res_hidden_states_tuple = res_hidden_states_tuple[-1] + if res_hidden_states_tuple is not None: + hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1) + + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + for resnet, attn in zip(self.resnets, self.attentions): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +# can potentially later be renamed to `No-feed-forward` attention +class KAttentionBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout: float = 0.0, + cross_attention_dim: Optional[int] = None, + attention_bias: bool = False, + upcast_attention: bool = False, + temb_channels: int = 768, # for ada_group_norm + add_self_attention: bool = False, + cross_attention_norm: Optional[str] = None, + group_size: int = 32, + ): + super().__init__() + self.add_self_attention = add_self_attention + + # 1. Self-Attn + if add_self_attention: + self.norm1 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size)) + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=None, + cross_attention_norm=None, + ) + + # 2. Cross-Attn + self.norm2 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size)) + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + cross_attention_norm=cross_attention_norm, + ) + + def _to_3d(self, hidden_states, height, weight): + return hidden_states.permute(0, 2, 3, 1).reshape(hidden_states.shape[0], height * weight, -1) + + def _to_4d(self, hidden_states, height, weight): + return hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, height, weight) + + def forward( + self, + hidden_states: torch.FloatTensor, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + # TODO: mark emb as non-optional (self.norm2 requires it). + # requires assessing impact of change to positional param interface. + emb: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + # 1. Self-Attention + if self.add_self_attention: + norm_hidden_states = self.norm1(hidden_states, emb) + + height, weight = norm_hidden_states.shape[2:] + norm_hidden_states = self._to_3d(norm_hidden_states, height, weight) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + attn_output = self._to_4d(attn_output, height, weight) + + hidden_states = attn_output + hidden_states + + # 2. Cross-Attention/None + norm_hidden_states = self.norm2(hidden_states, emb) + + height, weight = norm_hidden_states.shape[2:] + norm_hidden_states = self._to_3d(norm_hidden_states, height, weight) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask if encoder_hidden_states is None else encoder_attention_mask, + **cross_attention_kwargs, + ) + attn_output = self._to_4d(attn_output, height, weight) + + hidden_states = attn_output + hidden_states + + return hidden_states diff --git a/ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_condition.py b/ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_condition.py new file mode 100644 index 00000000..ce09d3bf --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/sdxl/unet_2d_condition.py @@ -0,0 +1,1259 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.utils.checkpoint +class FourierEmbedder(nn.Module): + def __init__(self, num_freqs=64, temperature=100): + super().__init__() + + self.num_freqs = num_freqs + self.temperature = temperature + + freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs) + freq_bands = freq_bands[None, None, None] + self.register_buffer("freq_bands", freq_bands, persistent=False) + + def __call__(self, x): + x = self.freq_bands * x.unsqueeze(-1) + return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 1, 3, 4, 2).reshape(*x.shape[:2], -1) + +class PositionNet(nn.Module): + def __init__(self, positive_len, out_dim, feature_type="text-only", fourier_freqs=8): + super().__init__() + self.positive_len = positive_len + self.out_dim = out_dim + + self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs) + self.position_dim = fourier_freqs * 2 * 4 # 2: sin/cos, 4: xyxy + + if isinstance(out_dim, tuple): + out_dim = out_dim[0] + + if feature_type == "text-only": + self.linears = nn.Sequential( + nn.Linear(self.positive_len + self.position_dim, 512), + nn.SiLU(), + nn.Linear(512, 512), + nn.SiLU(), + nn.Linear(512, out_dim), + ) + self.null_positive_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) + + elif feature_type == "text-image": + self.linears_text = nn.Sequential( + nn.Linear(self.positive_len + self.position_dim, 512), + nn.SiLU(), + nn.Linear(512, 512), + nn.SiLU(), + nn.Linear(512, out_dim), + ) + self.linears_image = nn.Sequential( + nn.Linear(self.positive_len + self.position_dim, 512), + nn.SiLU(), + nn.Linear(512, 512), + nn.SiLU(), + nn.Linear(512, out_dim), + ) + self.null_text_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) + self.null_image_feature = torch.nn.Parameter(torch.zeros([self.positive_len])) + + self.null_position_feature = torch.nn.Parameter(torch.zeros([self.position_dim])) + + def forward( + self, + boxes, + masks, + positive_embeddings=None, + phrases_masks=None, + image_masks=None, + phrases_embeddings=None, + image_embeddings=None, + ): + masks = masks.unsqueeze(-1) + + # embedding position (it may includes padding as placeholder) + xyxy_embedding = self.fourier_embedder(boxes) # B*N*4 -> B*N*C + + # learnable null embedding + xyxy_null = self.null_position_feature.view(1, 1, -1) + + # replace padding with learnable null embedding + xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null + + # positionet with text only information + if positive_embeddings is not None: + # learnable null embedding + positive_null = self.null_positive_feature.view(1, 1, -1) + + # replace padding with learnable null embedding + positive_embeddings = positive_embeddings * masks + (1 - masks) * positive_null + + objs = self.linears(torch.cat([positive_embeddings, xyxy_embedding], dim=-1)) + + # positionet with text and image infomation + else: + phrases_masks = phrases_masks.unsqueeze(-1) + image_masks = image_masks.unsqueeze(-1) + + # learnable null embedding + text_null = self.null_text_feature.view(1, 1, -1) + image_null = self.null_image_feature.view(1, 1, -1) + + # replace padding with learnable null embedding + phrases_embeddings = phrases_embeddings * phrases_masks + (1 - phrases_masks) * text_null + image_embeddings = image_embeddings * image_masks + (1 - image_masks) * image_null + + objs_text = self.linears_text(torch.cat([phrases_embeddings, xyxy_embedding], dim=-1)) + objs_image = self.linears_image(torch.cat([image_embeddings, xyxy_embedding], dim=-1)) + objs = torch.cat([objs_text, objs_image], dim=1) + + return objs + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import UNet2DConditionLoadersMixin +from diffusers.utils import BaseOutput, logging +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import ( + ADDED_KV_ATTENTION_PROCESSORS, + CROSS_ATTENTION_PROCESSORS, + AttentionProcessor, + AttnAddedKVProcessor, + AttnProcessor, +) +from diffusers.models.embeddings import ( + GaussianFourierProjection, + ImageHintTimeEmbedding, + ImageProjection, + ImageTimeEmbedding, + # PositionNet, + TextImageProjection, + TextImageTimeEmbedding, + TextTimeEmbedding, + TimestepEmbedding, + Timesteps, +) +from diffusers.models.modeling_utils import ModelMixin + +from .unet_2d_blocks import ( + UNetMidBlock2DCrossAttn, + UNetMidBlock2DSimpleCrossAttn, + get_down_block, + get_up_block, +) + +import time + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class UNet2DConditionOutput(BaseOutput): + """ + The output of [`UNet2DConditionModel`]. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. + """ + + sample: torch.FloatTensor = None + + +class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): + r""" + A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample + shaped output. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Parameters: + sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + Height and width of input/output sample. + in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample. + out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. + center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample. + flip_sin_to_cos (`bool`, *optional*, defaults to `False`): + Whether to flip the sin to cos in the time embedding. + freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. + down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): + The tuple of downsample blocks to use. + mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): + Block type for middle of UNet, it can be either `UNetMidBlock2DCrossAttn` or + `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): + The tuple of upsample blocks to use. + only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`): + Whether to include self-attention in the basic transformer blocks, see + [`~models.attention.BasicTransformerBlock`]. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + The tuple of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. + downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. + mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. + norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. + If `None`, normalization and activation layers is skipped in post-processing. + norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization. + cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + The dimension of the cross attention features. + transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1): + The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for + [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], + [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. + encoder_hid_dim (`int`, *optional*, defaults to None): + If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` + dimension to `cross_attention_dim`. + encoder_hid_dim_type (`str`, *optional*, defaults to `None`): + If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text + embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`. + attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads. + num_attention_heads (`int`, *optional*): + The number of attention heads. If not defined, defaults to `attention_head_dim` + resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config + for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`. + class_embed_type (`str`, *optional*, defaults to `None`): + The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, + `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. + addition_embed_type (`str`, *optional*, defaults to `None`): + Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or + "text". "text" will use the `TextTimeEmbedding` layer. + addition_time_embed_dim: (`int`, *optional*, defaults to `None`): + Dimension for the timestep embeddings. + num_class_embeds (`int`, *optional*, defaults to `None`): + Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing + class conditioning with `class_embed_type` equal to `None`. + time_embedding_type (`str`, *optional*, defaults to `positional`): + The type of position embedding to use for timesteps. Choose from `positional` or `fourier`. + time_embedding_dim (`int`, *optional*, defaults to `None`): + An optional override for the dimension of the projected time embedding. + time_embedding_act_fn (`str`, *optional*, defaults to `None`): + Optional activation function to use only once on the time embeddings before they are passed to the rest of + the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`. + timestep_post_act (`str`, *optional*, defaults to `None`): + The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`. + time_cond_proj_dim (`int`, *optional*, defaults to `None`): + The dimension of `cond_proj` layer in the timestep embedding. + conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. + conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer. + projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when + `class_embed_type="projection"`. Required when `class_embed_type="projection"`. + class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time + embeddings with the class embeddings. + mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`): + Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If + `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the + `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False` + otherwise. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + center_input_sample: bool = False, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "DownBlock2D", + ), + mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", + up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: Union[int, Tuple[int]] = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + dropout: float = 0.0, + act_fn: str = "silu", + norm_num_groups: Optional[int] = 32, + norm_eps: float = 1e-5, + cross_attention_dim: Union[int, Tuple[int]] = 1280, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + encoder_hid_dim: Optional[int] = None, + encoder_hid_dim_type: Optional[str] = None, + attention_head_dim: Union[int, Tuple[int]] = 8, + num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + addition_embed_type: Optional[str] = None, + addition_time_embed_dim: Optional[int] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: int = 1.0, + time_embedding_type: str = "positional", + time_embedding_dim: Optional[int] = None, + time_embedding_act_fn: Optional[str] = None, + timestep_post_act: Optional[str] = None, + time_cond_proj_dim: Optional[int] = None, + conv_in_kernel: int = 3, + conv_out_kernel: int = 3, + projection_class_embeddings_input_dim: Optional[int] = None, + attention_type: str = "default", + class_embeddings_concat: bool = False, + mid_block_only_cross_attention: Optional[bool] = None, + cross_attention_norm: Optional[str] = None, + addition_embed_type_num_heads=64, + ): + super().__init__() + + self.sample_size = sample_size + + if num_attention_heads is not None: + raise ValueError( + "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." + ) + + # If `num_attention_heads` is not defined (which is the case for most models) + # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. + # The reason for this behavior is to correct for incorrectly named variables that were introduced + # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 + # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking + # which is why we correct for the naming here. + num_attention_heads = num_attention_heads or attention_head_dim + + # Check inputs + if len(down_block_types) != len(up_block_types): + raise ValueError( + f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." + ) + + if len(block_out_channels) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}." + ) + + if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." + ) + + # input + conv_in_padding = (conv_in_kernel - 1) // 2 + self.conv_in = nn.Conv2d( + in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding + ) + + # time + if time_embedding_type == "fourier": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 + if time_embed_dim % 2 != 0: + raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.") + self.time_proj = GaussianFourierProjection( + time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos + ) + timestep_input_dim = time_embed_dim + elif time_embedding_type == "positional": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 + + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) + timestep_input_dim = block_out_channels[0] + else: + raise ValueError( + f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." + ) + + self.time_embedding = TimestepEmbedding( + timestep_input_dim, + time_embed_dim, + act_fn=act_fn, + post_act_fn=timestep_post_act, + cond_proj_dim=time_cond_proj_dim, + ) + + if encoder_hid_dim_type is None and encoder_hid_dim is not None: + encoder_hid_dim_type = "text_proj" + self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type) + logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.") + + if encoder_hid_dim is None and encoder_hid_dim_type is not None: + raise ValueError( + f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}." + ) + + if encoder_hid_dim_type == "text_proj": + self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim) + elif encoder_hid_dim_type == "text_image_proj": + # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)` + self.encoder_hid_proj = TextImageProjection( + text_embed_dim=encoder_hid_dim, + image_embed_dim=cross_attention_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 + self.encoder_hid_proj = ImageProjection( + image_embed_dim=encoder_hid_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type is not None: + raise ValueError( + f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'." + ) + else: + self.encoder_hid_proj = None + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) + elif class_embed_type == "timestep": + self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn) + elif class_embed_type == "identity": + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + elif class_embed_type == "projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" + ) + # The projection `class_embed_type` is the same as the timestep `class_embed_type` except + # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings + # 2. it projects from an arbitrary input dimension. + # + # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. + # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. + # As a result, `TimestepEmbedding` can be passed arbitrary vectors. + self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) + elif class_embed_type == "simple_projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set" + ) + self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim) + else: + self.class_embedding = None + + if addition_embed_type == "text": + if encoder_hid_dim is not None: + text_time_embedding_from_dim = encoder_hid_dim + else: + text_time_embedding_from_dim = cross_attention_dim + + self.add_embedding = TextTimeEmbedding( + text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads + ) + elif addition_embed_type == "text_image": + # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)` + self.add_embedding = TextImageTimeEmbedding( + text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim + ) + elif addition_embed_type == "text_time": + self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) + self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) + elif addition_embed_type == "image": + # Kandinsky 2.2 + self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) + elif addition_embed_type == "image_hint": + # Kandinsky 2.2 ControlNet + self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim) + elif addition_embed_type is not None: + raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.") + + if time_embedding_act_fn is None: + self.time_embed_act = None + else: + self.time_embed_act = get_activation(time_embedding_act_fn) + + self.down_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = only_cross_attention + + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = False + + if isinstance(num_attention_heads, int): + num_attention_heads = (num_attention_heads,) * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + if isinstance(cross_attention_dim, int): + cross_attention_dim = (cross_attention_dim,) * len(down_block_types) + + if isinstance(layers_per_block, int): + layers_per_block = [layers_per_block] * len(down_block_types) + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) + + if class_embeddings_concat: + # The time embeddings are concatenated with the class embeddings. The dimension of the + # time embeddings passed to the down, middle, and up blocks is twice the dimension of the + # regular time embeddings + blocks_time_embed_dim = time_embed_dim * 2 + else: + blocks_time_embed_dim = time_embed_dim + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block[i], + transformer_layers_per_block=transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + temb_channels=blocks_time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim[i], + num_attention_heads=num_attention_heads[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + dropout=dropout, + ) + self.down_blocks.append(down_block) + + # mid + if mid_block_type == "UNetMidBlock2DCrossAttn": + self.mid_block = UNetMidBlock2DCrossAttn( + transformer_layers_per_block=transformer_layers_per_block[-1], + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + dropout=dropout, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim[-1], + num_attention_heads=num_attention_heads[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn": + self.mid_block = UNetMidBlock2DSimpleCrossAttn( + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + dropout=dropout, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + cross_attention_dim=cross_attention_dim[-1], + attention_head_dim=attention_head_dim[-1], + resnet_groups=norm_num_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + only_cross_attention=mid_block_only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif mid_block_type is None: + self.mid_block = None + else: + raise ValueError(f"unknown mid_block_type : {mid_block_type}") + + # count how many layers upsample the images + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_num_attention_heads = list(reversed(num_attention_heads)) + reversed_layers_per_block = list(reversed(layers_per_block)) + reversed_cross_attention_dim = list(reversed(cross_attention_dim)) + reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) + only_cross_attention = list(reversed(only_cross_attention)) + + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=reversed_layers_per_block[i] + 1, + transformer_layers_per_block=reversed_transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=blocks_time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=reversed_cross_attention_dim[i], + num_attention_heads=reversed_num_attention_heads[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel, + dropout=dropout, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_num_groups is not None: + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps + ) + + self.conv_act = get_activation(act_fn) + + else: + self.conv_norm_out = None + self.conv_act = None + + conv_out_padding = (conv_out_kernel - 1) // 2 + self.conv_out = nn.Conv2d( + block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding + ) + + if attention_type in ["gated", "gated-text-image"]: + positive_len = 768 + if isinstance(cross_attention_dim, int): + positive_len = cross_attention_dim + elif isinstance(cross_attention_dim, tuple) or isinstance(cross_attention_dim, list): + positive_len = cross_attention_dim[0] + + feature_type = "text-only" if attention_type == "gated" else "text-image" + self.position_net = PositionNet( + positive_len=positive_len, out_dim=cross_attention_dim, feature_type=feature_type + ) + + @property + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attn_processor( + self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], _remove_lora=False + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor, _remove_lora=_remove_lora) + else: + module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnAddedKVProcessor() + elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor, _remove_lora=True) + + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module splits the input tensor in slices to compute attention in + several steps. This is useful for saving some memory in exchange for a small decrease in speed. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If + `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): + if hasattr(module, "set_attention_slice"): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_sliceable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_sliceable_dims(module) + + num_sliceable_layers = len(sliceable_head_dims) + + if slice_size == "auto": + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == "max": + # make smallest slice possible + slice_size = num_sliceable_layers * [1] + + slice_size = num_sliceable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" + f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError(f"size {size} has to be smaller or equal to {dim}.") + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]): + if hasattr(module, "set_attention_slice"): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + class_labels: Optional[torch.Tensor] = None, + timestep_cond: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + mid_block_additional_residual: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + quick_replicate: bool = False, + replicate_prv_feature: Optional[List[torch.Tensor]] = None, + cache_layer_id: Optional[int] = None, + cache_block_id: Optional[int] = None, + return_dict: bool = True, + ) -> Union[UNet2DConditionOutput, Tuple]: + r""" + The [`UNet2DConditionModel`] forward method. + + Args: + sample (`torch.FloatTensor`): + The noisy input tensor with the following shape `(batch, channel, height, width)`. + timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. + encoder_hidden_states (`torch.FloatTensor`): + The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. + encoder_attention_mask (`torch.Tensor`): + A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If + `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, + which adds large negative values to the attention scores corresponding to "discard" tokens. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + + Returns: + [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: + If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise + a `tuple` is returned where the first element is the sample tensor. + """ + # By default samples have to be AT least a multiple of the overall upsampling factor. + # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). + # However, the upsampling interpolation output size can be forced to fit any upsampling size + # on the fly if necessary. + default_overall_up_factor = 2**self.num_upsamplers + + # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` + forward_upsample_size = False + upsample_size = None + + if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): + logger.info("Forward upsample size to force interpolation output size.") + forward_upsample_size = True + + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension + # expects mask of shape: + # [batch, key_tokens] + # adds singleton query_tokens dimension: + # [batch, 1, key_tokens] + # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: + # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) + # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) + if attention_mask is not None: + # assume that mask is expressed as: + # (1 = keep, 0 = discard) + # convert mask into a bias that can be added to attention scores: + # (keep = +0, discard = -10000.0) + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # convert encoder_attention_mask to a bias the same way we do for attention_mask + if encoder_attention_mask is not None: + encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # 0. center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb, timestep_cond) + aug_emb = None + + if self.class_embedding is not None: + if class_labels is None: + raise ValueError("class_labels should be provided when num_class_embeds > 0") + + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # there might be better ways to encapsulate this. + class_labels = class_labels.to(dtype=sample.dtype) + + class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype) + + if self.config.class_embeddings_concat: + emb = torch.cat([emb, class_emb], dim=-1) + else: + emb = emb + class_emb + + if self.config.addition_embed_type == "text": + aug_emb = self.add_embedding(encoder_hidden_states) + elif self.config.addition_embed_type == "text_image": + # Kandinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + + image_embs = added_cond_kwargs.get("image_embeds") + text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states) + aug_emb = self.add_embedding(text_embs, image_embs) + elif self.config.addition_embed_type == "text_time": + # SDXL - style + if "text_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" + ) + text_embeds = added_cond_kwargs.get("text_embeds") + if "time_ids" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" + ) + time_ids = added_cond_kwargs.get("time_ids") + time_embeds = self.add_time_proj(time_ids.flatten()) + time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) + + add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) + add_embeds = add_embeds.to(emb.dtype) + aug_emb = self.add_embedding(add_embeds) + elif self.config.addition_embed_type == "image": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + aug_emb = self.add_embedding(image_embs) + elif self.config.addition_embed_type == "image_hint": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + hint = added_cond_kwargs.get("hint") + aug_emb, hint = self.add_embedding(image_embs, hint) + sample = torch.cat([sample, hint], dim=1) + + emb = emb + aug_emb if aug_emb is not None else emb + + if self.time_embed_act is not None: + emb = self.time_embed_act(emb) + + if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj": + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj": + # Kadinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds) + elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(image_embeds) + # 2. pre-process + sample = self.conv_in(sample) + + # 2.5 GLIGEN position net + if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None: + cross_attention_kwargs = cross_attention_kwargs.copy() + gligen_args = cross_attention_kwargs.pop("gligen") + cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)} + + # 3. down + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None + is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None + + down_block_res_samples = (sample,) + if quick_replicate and replicate_prv_feature is not None: + # Down + for i, downsample_block in enumerate(self.down_blocks): + if i > cache_layer_id: + break + + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + # For t2i-adapter CrossAttnDownBlock2D + additional_residuals = {} + if is_adapter and len(down_block_additional_residuals) > 0: + additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0) + + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + exist_block_number=cache_block_id if i == cache_layer_id else None, + **additional_residuals, + ) + else: + sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale, exist_block_number=cache_block_id if i == cache_layer_id else None,) + + if is_adapter and len(down_block_additional_residuals) > 0: + sample += down_block_additional_residuals.pop(0) + + down_block_res_samples += res_samples + + # No Middle + # Up + #print("down_block_res_samples:", [res_sample.shape for res_sample in down_block_res_samples]) + sample = replicate_prv_feature + max_block_depth = len(self.down_blocks[cache_layer_id].attentions) if hasattr(self.down_blocks[cache_layer_id], "attentions") else len(self.down_blocks[cache_layer_id].resnets) + if cache_block_id == max_block_depth : + cache_block_id = 0 + cache_layer_id += 1 + else: + cache_block_id += 1 + + for i, upsample_block in enumerate(self.up_blocks): + if i < len(self.up_blocks) - 1 - cache_layer_id: + continue + + if i == len(self.up_blocks) - 1 - cache_layer_id: + trunc_upsample_block = cache_block_id + 1 + else: + trunc_upsample_block = len(upsample_block.resnets) + + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-trunc_upsample_block:] + down_block_res_samples = down_block_res_samples[: -trunc_upsample_block] + + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + #print(sample.shape, [res_sample.shape for res_sample in res_samples]) + sample, _ = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + enter_block_number=cache_block_id if i == len(self.up_blocks) - 1 - cache_layer_id else None, + ) + else: + sample, _ = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size, + scale=lora_scale, + enter_block_number=cache_block_id if i == len(self.up_blocks) - 1 - cache_layer_id else None, + ) + + prv_f = replicate_prv_feature + else: + for i, downsample_block in enumerate(self.down_blocks): + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + # For t2i-adapter CrossAttnDownBlock2D + additional_residuals = {} + if is_adapter and len(down_block_additional_residuals) > 0: + additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0) + + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + **additional_residuals, + ) + else: + sample, res_samples = downsample_block(hidden_states=sample, temb=emb, scale=lora_scale) + + if is_adapter and len(down_block_additional_residuals) > 0: + sample += down_block_additional_residuals.pop(0) + + down_block_res_samples += res_samples + + if is_controlnet: + new_down_block_res_samples = () + + for down_block_res_sample, down_block_additional_residual in zip( + down_block_res_samples, down_block_additional_residuals + ): + down_block_res_sample = down_block_res_sample + down_block_additional_residual + new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) + + down_block_res_samples = new_down_block_res_samples + + # 4. mid + if self.mid_block is not None: + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + # To support T2I-Adapter-XL + if ( + is_adapter + and len(down_block_additional_residuals) > 0 + and sample.shape == down_block_additional_residuals[0].shape + ): + sample += down_block_additional_residuals.pop(0) + + if is_controlnet: + sample = sample + mid_block_additional_residual + + # 5. up + if cache_block_id is not None: + max_block_depth = len(self.down_blocks[cache_layer_id].attentions) if hasattr(self.down_blocks[cache_layer_id], "attentions") else len(self.down_blocks[cache_layer_id].resnets) + if cache_block_id == max_block_depth: + cache_block_id = 0 + cache_layer_id += 1 + else: + cache_block_id += 1 + #print("down_block_res_samples:", [res_sample.shape for res_sample in down_block_res_samples]) + #print(cache_block_id, cache_layer_id) + prv_f = None + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] + #print(sample.shape, [res_sample.shape for res_sample in res_samples]) + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + sample, current_record_f = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample, current_record_f = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size, + scale=lora_scale, + ) + + #print(cache_layer_id, current_record_f is None, i == len(self.up_blocks) - cache_layer_id - 1) + #print("Append prv_feature with shape:", sample.shape) + if cache_layer_id is not None and current_record_f is not None and i == len(self.up_blocks) - cache_layer_id - 1: + prv_f = current_record_f[-cache_block_id-1] + + # 6. post-process + if self.conv_norm_out: + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + if not return_dict: + return (sample, prv_f,) + return UNet2DConditionOutput(sample=sample) diff --git a/ixformer_sdk/contrib/DeepCache/svd/__init__.py b/ixformer_sdk/contrib/DeepCache/svd/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/DeepCache/svd/pipeline_stable_video_diffusion.py b/ixformer_sdk/contrib/DeepCache/svd/pipeline_stable_video_diffusion.py new file mode 100644 index 00000000..b318a49d --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/svd/pipeline_stable_video_diffusion.py @@ -0,0 +1,659 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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 inspect +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Union + +import numpy as np +import PIL.Image +import torch +from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection + +from diffusers.image_processor import VaeImageProcessor +from diffusers.models import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel +from diffusers.schedulers import EulerDiscreteScheduler +from diffusers.utils import BaseOutput, logging +from diffusers.utils.torch_utils import randn_tensor +from .pipeline_utils import DiffusionPipeline + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _append_dims(x, target_dims): + """Appends dimensions to the end of a tensor until it has target_dims dimensions.""" + dims_to_append = target_dims - x.ndim + if dims_to_append < 0: + raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less") + return x[(...,) + (None,) * dims_to_append] + + +def tensor2vid(video: torch.Tensor, processor, output_type="np"): + # Based on: + # https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78 + + batch_size, channels, num_frames, height, width = video.shape + outputs = [] + for batch_idx in range(batch_size): + batch_vid = video[batch_idx].permute(1, 0, 2, 3) + batch_output = processor.postprocess(batch_vid, output_type) + + outputs.append(batch_output) + + return outputs + + +@dataclass +class StableVideoDiffusionPipelineOutput(BaseOutput): + r""" + Output class for zero-shot text-to-video pipeline. + + Args: + frames (`[List[PIL.Image.Image]`, `np.ndarray`]): + List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + num_channels)`. + """ + + frames: Union[List[PIL.Image.Image], np.ndarray] + + +class StableVideoDiffusionPipeline(DiffusionPipeline): + r""" + Pipeline to generate video from an input image using Stable Video Diffusion. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. + image_encoder ([`~transformers.CLIPVisionModelWithProjection`]): + Frozen CLIP image-encoder ([laion/CLIP-ViT-H-14-laion2B-s32B-b79K](https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K)). + unet ([`UNetSpatioTemporalConditionModel`]):cache_interval=5, cache_branch=0, + A `UNetSpatioTemporalConditionModel` to denoise the encoded image latents. + scheduler ([`EulerDiscreteScheduler`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. + feature_extractor ([`~transformers.CLIPImageProcessor`]): + A `CLIPImageProcessor` to extract features from generated images. + """ + + model_cpu_offload_seq = "image_encoder->unet->vae" + _callback_tensor_inputs = ["latents"] + + def __init__( + self, + vae: AutoencoderKLTemporalDecoder, + image_encoder: CLIPVisionModelWithProjection, + unet: UNetSpatioTemporalConditionModel, + scheduler: EulerDiscreteScheduler, + feature_extractor: CLIPImageProcessor, + ): + super().__init__() + + self.register_modules( + vae=vae, + image_encoder=image_encoder, + unet=unet, + scheduler=scheduler, + feature_extractor=feature_extractor, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + + def _encode_image(self, image, device, num_videos_per_prompt, do_classifier_free_guidance): + dtype = next(self.image_encoder.parameters()).dtype + + if not isinstance(image, torch.Tensor): + image = self.image_processor.pil_to_numpy(image) + image = self.image_processor.numpy_to_pt(image) + + # We normalize the image before resizing to match with the original implementation. + # Then we unnormalize it after resizing. + image = image * 2.0 - 1.0 + image = _resize_with_antialiasing(image, (224, 224)) + image = (image + 1.0) / 2.0 + + # Normalize the image with for CLIP input + image = self.feature_extractor( + images=image, + do_normalize=True, + do_center_crop=False, + do_resize=False, + do_rescale=False, + return_tensors="pt", + ).pixel_values + + image = image.to(device=device, dtype=dtype) + image_embeddings = self.image_encoder(image).image_embeds + image_embeddings = image_embeddings.unsqueeze(1) + + # duplicate image embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = image_embeddings.shape + image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1) + image_embeddings = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + negative_image_embeddings = torch.zeros_like(image_embeddings) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_embeddings = torch.cat([negative_image_embeddings, image_embeddings]) + + return image_embeddings + + def _encode_vae_image( + self, + image: torch.Tensor, + device, + num_videos_per_prompt, + do_classifier_free_guidance, + ): + image = image.to(device=device) + image_latents = self.vae.encode(image).latent_dist.mode() + + if do_classifier_free_guidance: + negative_image_latents = torch.zeros_like(image_latents) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_latents = torch.cat([negative_image_latents, image_latents]) + + # duplicate image_latents for each generation per prompt, using mps friendly method + image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1) + + return image_latents + + def _get_add_time_ids( + self, + fps, + motion_bucket_id, + noise_aug_strength, + dtype, + batch_size, + num_videos_per_prompt, + do_classifier_free_guidance, + ): + add_time_ids = [fps, motion_bucket_id, noise_aug_strength] + + passed_add_embed_dim = self.unet.config.addition_time_embed_dim * len(add_time_ids) + expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features + + if expected_add_embed_dim != passed_add_embed_dim: + raise ValueError( + f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." + ) + + add_time_ids = torch.tensor([add_time_ids], dtype=dtype) + add_time_ids = add_time_ids.repeat(batch_size * num_videos_per_prompt, 1) + + if do_classifier_free_guidance: + add_time_ids = torch.cat([add_time_ids, add_time_ids]) + + return add_time_ids + + def decode_latents(self, latents, num_frames, decode_chunk_size=14): + # [batch, frames, channels, height, width] -> [batch*frames, channels, height, width] + latents = latents.flatten(0, 1) + + latents = 1 / self.vae.config.scaling_factor * latents + + accepts_num_frames = "num_frames" in set(inspect.signature(self.vae.forward).parameters.keys()) + + # decode decode_chunk_size frames at a time to avoid OOM + frames = [] + for i in range(0, latents.shape[0], decode_chunk_size): + num_frames_in = latents[i : i + decode_chunk_size].shape[0] + decode_kwargs = {} + if accepts_num_frames: + # we only pass num_frames_in if it's expected + decode_kwargs["num_frames"] = num_frames_in + + frame = self.vae.decode(latents[i : i + decode_chunk_size], **decode_kwargs).sample + frames.append(frame) + frames = torch.cat(frames, dim=0) + + # [batch*frames, channels, height, width] -> [batch, channels, frames, height, width] + frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4) + + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 + frames = frames.float() + return frames + + def check_inputs(self, image, height, width): + if ( + not isinstance(image, torch.Tensor) + and not isinstance(image, PIL.Image.Image) + and not isinstance(image, list) + ): + raise ValueError( + "`image` has to be of type `torch.FloatTensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is" + f" {type(image)}" + ) + + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + def prepare_latents( + self, + batch_size, + num_frames, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + shape = ( + batch_size, + num_frames, + num_channels_latents // 2, + height // self.vae_scale_factor, + width // self.vae_scale_factor, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None + + @property + def num_timesteps(self): + return self._num_timesteps + + @torch.no_grad() + def __call__( + self, + image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.FloatTensor], + height: int = 576, + width: int = 1024, + num_frames: Optional[int] = None, + num_inference_steps: int = 25, + min_guidance_scale: float = 1.0, + max_guidance_scale: float = 3.0, + fps: int = 7, + motion_bucket_id: int = 127, + noise_aug_strength: int = 0.02, + decode_chunk_size: Optional[int] = None, + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + cache_interval: Optional[int] = 1, + cache_branch: Optional[int] = None, + return_dict: bool = True, + ): + r""" + The call function to the pipeline for generation. + + Args: + image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`): + Image or images to guide image generation. If you provide a tensor, it needs to be compatible with + [`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json). + height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`): + The width in pixels of the generated image. + num_frames (`int`, *optional*): + The number of video frames to generate. Defaults to 14 for `stable-video-diffusion-img2vid` and to 25 for `stable-video-diffusion-img2vid-xt` + num_inference_steps (`int`, *optional*, defaults to 25): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. This parameter is modulated by `strength`. + min_guidance_scale (`float`, *optional*, defaults to 1.0): + The minimum guidance scale. Used for the classifier free guidance with first frame. + max_guidance_scale (`float`, *optional*, defaults to 3.0): + The maximum guidance scale. Used for the classifier free guidance with last frame. + fps (`int`, *optional*, defaults to 7): + Frames per second. The rate at which the generated images shall be exported to a video after generation. + Note that Stable Diffusion Video's UNet was micro-conditioned on fps-1 during training. + motion_bucket_id (`int`, *optional*, defaults to 127): + The motion bucket ID. Used as conditioning for the generation. The higher the number the more motion will be in the video. + noise_aug_strength (`int`, *optional*, defaults to 0.02): + The amount of noise added to the init image, the higher it is the less the video will look like the init image. Increase it for more motion. + decode_chunk_size (`int`, *optional*): + The number of frames to decode at a time. The higher the chunk size, the higher the temporal consistency + between frames, but also the higher the memory consumption. By default, the decoder will decode all frames at once + for maximal quality. Reduce `decode_chunk_size` to reduce memory usage. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a + plain tuple. + + Returns: + [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] is returned, + otherwise a `tuple` is returned where the first element is a list of list with the generated frames. + + Examples: + + ```py + from diffusers import StableVideoDiffusionPipeline + from diffusers.utils import load_image, export_to_video + + pipe = StableVideoDiffusionPipeline.from_pretrained("stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16") + pipe.to("cuda") + + image = load_image("https://lh3.googleusercontent.com/y-iFOHfLTwkuQSUegpwDdgKmOjRSTvPxat63dQLB25xkTs4lhIbRUFeNBWZzYf370g=s1200") + image = image.resize((1024, 576)) + + frames = pipe(image, num_frames=25, decode_chunk_size=8).frames[0] + export_to_video(frames, "generated.mp4", fps=7) + ``` + """ + # 0. Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + num_frames = num_frames if num_frames is not None else self.unet.config.num_frames + decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else num_frames + + # 1. Check inputs. Raise error if not correct + self.check_inputs(image, height, width) + + # 2. Define call parameters + if isinstance(image, PIL.Image.Image): + batch_size = 1 + elif isinstance(image, list): + batch_size = len(image) + else: + batch_size = image.shape[0] + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = max_guidance_scale > 1.0 + + # 3. Encode input image + image_embeddings = self._encode_image(image, device, num_videos_per_prompt, do_classifier_free_guidance) + + # NOTE: Stable Diffusion Video was conditioned on fps - 1, which + # is why it is reduced here. + # See: https://github.com/Stability-AI/generative-models/blob/ed0997173f98eaf8f4edf7ba5fe8f15c6b877fd3/scripts/sampling/simple_video_sample.py#L188 + fps = fps - 1 + + # 4. Encode input image using VAE + image = self.image_processor.preprocess(image, height=height, width=width) + noise = randn_tensor(image.shape, generator=generator, device=image.device, dtype=image.dtype) + image = image + noise_aug_strength * noise + + needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast + if needs_upcasting: + self.vae.to(dtype=torch.float32) + + image_latents = self._encode_vae_image(image, device, num_videos_per_prompt, do_classifier_free_guidance) + image_latents = image_latents.to(image_embeddings.dtype) + + # cast back to fp16 if needed + if needs_upcasting: + self.vae.to(dtype=torch.float16) + + # Repeat the image latents for each frame so we can concatenate them with the noise + # image_latents [batch, channels, height, width] ->[batch, num_frames, channels, height, width] + image_latents = image_latents.unsqueeze(1).repeat(1, num_frames, 1, 1, 1) + + # 5. Get Added Time IDs + added_time_ids = self._get_add_time_ids( + fps, + motion_bucket_id, + noise_aug_strength, + image_embeddings.dtype, + batch_size, + num_videos_per_prompt, + do_classifier_free_guidance, + ) + added_time_ids = added_time_ids.to(device) + + # 4. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # 5. Prepare latent variables + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_frames, + num_channels_latents, + height, + width, + image_embeddings.dtype, + device, + generator, + latents, + ) + + # 7. Prepare guidance scale + guidance_scale = torch.linspace(min_guidance_scale, max_guidance_scale, num_frames).unsqueeze(0) + guidance_scale = guidance_scale.to(device, latents.dtype) + guidance_scale = guidance_scale.repeat(batch_size * num_videos_per_prompt, 1) + guidance_scale = _append_dims(guidance_scale, latents.ndim) + + self._guidance_scale = guidance_scale + + cache_features = None + interval_seq = list(range(0, num_inference_steps, cache_interval)) + interval_seq = sorted(interval_seq) + + # 8. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + # Concatenate image_latents over channels dimention + latent_model_input = torch.cat([latent_model_input, image_latents], dim=2) + + if i in interval_seq: + cache_features = None + + # predict the noise residual + noise_pred, cache_features = self.unet( + latent_model_input, + t, + encoder_hidden_states=image_embeddings, + added_time_ids=added_time_ids, + cache_features=cache_features, + cache_branch=cache_branch, + return_dict=False, + ) + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents).prev_sample + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if not output_type == "latent": + # cast back to fp16 if needed + if needs_upcasting: + self.vae.to(dtype=torch.float16) + frames = self.decode_latents(latents, num_frames, decode_chunk_size) + frames = tensor2vid(frames, self.image_processor, output_type=output_type) + else: + frames = latents + + self.maybe_free_model_hooks() + + if not return_dict: + return frames + + return StableVideoDiffusionPipelineOutput(frames=frames) + + +# resizing utils +# TODO: clean up later +def _resize_with_antialiasing(input, size, interpolation="bicubic", align_corners=True): + h, w = input.shape[-2:] + factors = (h / size[0], w / size[1]) + + # First, we have to determine sigma + # Taken from skimage: https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/transform/_warps.py#L171 + sigmas = ( + max((factors[0] - 1.0) / 2.0, 0.001), + max((factors[1] - 1.0) / 2.0, 0.001), + ) + + # Now kernel size. Good results are for 3 sigma, but that is kind of slow. Pillow uses 1 sigma + # https://github.com/python-pillow/Pillow/blob/master/src/libImaging/Resample.c#L206 + # But they do it in the 2 passes, which gives better results. Let's try 2 sigmas for now + ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3)) + + # Make sure it is odd + if (ks[0] % 2) == 0: + ks = ks[0] + 1, ks[1] + + if (ks[1] % 2) == 0: + ks = ks[0], ks[1] + 1 + + input = _gaussian_blur2d(input, ks, sigmas) + + output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners) + return output + + +def _compute_padding(kernel_size): + """Compute padding tuple.""" + # 4 or 6 ints: (padding_left, padding_right,padding_top,padding_bottom) + # https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad + if len(kernel_size) < 2: + raise AssertionError(kernel_size) + computed = [k - 1 for k in kernel_size] + + # for even kernels we need to do asymmetric padding :( + out_padding = 2 * len(kernel_size) * [0] + + for i in range(len(kernel_size)): + computed_tmp = computed[-(i + 1)] + + pad_front = computed_tmp // 2 + pad_rear = computed_tmp - pad_front + + out_padding[2 * i + 0] = pad_front + out_padding[2 * i + 1] = pad_rear + + return out_padding + + +def _filter2d(input, kernel): + # prepare kernel + b, c, h, w = input.shape + tmp_kernel = kernel[:, None, ...].to(device=input.device, dtype=input.dtype) + + tmp_kernel = tmp_kernel.expand(-1, c, -1, -1) + + height, width = tmp_kernel.shape[-2:] + + padding_shape: list[int] = _compute_padding([height, width]) + input = torch.nn.functional.pad(input, padding_shape, mode="reflect") + + # kernel and input tensor reshape to align element-wise or batch-wise params + tmp_kernel = tmp_kernel.reshape(-1, 1, height, width) + input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1)) + + # convolve the tensor with the kernel. + output = torch.nn.functional.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1) + + out = output.view(b, c, h, w) + return out + + +def _gaussian(window_size: int, sigma): + if isinstance(sigma, float): + sigma = torch.tensor([[sigma]]) + + batch_size = sigma.shape[0] + + x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1) + + if window_size % 2 == 0: + x = x + 0.5 + + gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0))) + + return gauss / gauss.sum(-1, keepdim=True) + + +def _gaussian_blur2d(input, kernel_size, sigma): + if isinstance(sigma, tuple): + sigma = torch.tensor([sigma], dtype=input.dtype) + else: + sigma = sigma.to(dtype=input.dtype) + + ky, kx = int(kernel_size[0]), int(kernel_size[1]) + bs = sigma.shape[0] + kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1)) + kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1)) + out_x = _filter2d(input, kernel_x[..., None, :]) + out = _filter2d(out_x, kernel_y[..., None]) + + return out diff --git a/ixformer_sdk/contrib/DeepCache/svd/pipeline_utils.py b/ixformer_sdk/contrib/DeepCache/svd/pipeline_utils.py new file mode 100644 index 00000000..1d40e6d4 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/svd/pipeline_utils.py @@ -0,0 +1,2108 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. +# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +# +# 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 fnmatch +import importlib +import inspect +import os +import re +import sys +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import PIL.Image +import torch +from huggingface_hub import ModelCard, create_repo, hf_hub_download, model_info, snapshot_download +from packaging import version +from requests.exceptions import HTTPError +from tqdm.auto import tqdm + +from diffusers import __version__ +from diffusers.configuration_utils import ConfigMixin +from diffusers.models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT +from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME +from diffusers.utils import ( + CONFIG_NAME, + DEPRECATED_REVISION_ARGS, + # DIFFUSERS_CACHE, + # HF_HUB_OFFLINE, + SAFETENSORS_WEIGHTS_NAME, + WEIGHTS_NAME, + BaseOutput, + deprecate, + get_class_from_dynamic_module, + is_accelerate_available, + is_accelerate_version, + is_peft_available, + is_torch_version, + is_transformers_available, + logging, + numpy_to_pil, +) +from diffusers.utils.torch_utils import is_compiled_module + +from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE +DIFFUSERS_CACHE=HUGGINGFACE_HUB_CACHE +ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} +HF_HUB_OFFLINE = os.getenv("HF_HUB_OFFLINE", "").upper() in ENV_VARS_TRUE_VALUES + +if is_transformers_available(): + import transformers + from transformers import PreTrainedModel + from transformers.utils import FLAX_WEIGHTS_NAME as TRANSFORMERS_FLAX_WEIGHTS_NAME + from transformers.utils import SAFE_WEIGHTS_NAME as TRANSFORMERS_SAFE_WEIGHTS_NAME + from transformers.utils import WEIGHTS_NAME as TRANSFORMERS_WEIGHTS_NAME + +from diffusers.utils import FLAX_WEIGHTS_NAME, ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, PushToHubMixin + + +if is_accelerate_available(): + import accelerate + + +INDEX_FILE = "diffusion_pytorch_model.bin" +CUSTOM_PIPELINE_FILE_NAME = "pipeline.py" +DUMMY_MODULES_FOLDER = "diffusers.utils" +TRANSFORMERS_DUMMY_MODULES_FOLDER = "transformers.utils" +CONNECTED_PIPES_KEYS = ["prior"] + + +logger = logging.get_logger(__name__) + + +LOADABLE_CLASSES = { + "diffusers": { + "ModelMixin": ["save_pretrained", "from_pretrained"], + "SchedulerMixin": ["save_pretrained", "from_pretrained"], + "DiffusionPipeline": ["save_pretrained", "from_pretrained"], + "OnnxRuntimeModel": ["save_pretrained", "from_pretrained"], + }, + "transformers": { + "PreTrainedTokenizer": ["save_pretrained", "from_pretrained"], + "PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"], + "PreTrainedModel": ["save_pretrained", "from_pretrained"], + "FeatureExtractionMixin": ["save_pretrained", "from_pretrained"], + "ProcessorMixin": ["save_pretrained", "from_pretrained"], + "ImageProcessingMixin": ["save_pretrained", "from_pretrained"], + }, + "onnxruntime.training": { + "ORTModule": ["save_pretrained", "from_pretrained"], + }, +} + +ALL_IMPORTABLE_CLASSES = {} +for library in LOADABLE_CLASSES: + ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library]) + + +@dataclass +class ImagePipelineOutput(BaseOutput): + """ + Output class for image pipelines. + + Args: + images (`List[PIL.Image.Image]` or `np.ndarray`) + List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width, + num_channels)`. + """ + + images: Union[List[PIL.Image.Image], np.ndarray] + + +@dataclass +class AudioPipelineOutput(BaseOutput): + """ + Output class for audio pipelines. + + Args: + audios (`np.ndarray`) + List of denoised audio samples of a NumPy array of shape `(batch_size, num_channels, sample_rate)`. + """ + + audios: np.ndarray + + +def is_safetensors_compatible(filenames, variant=None, passed_components=None) -> bool: + """ + Checking for safetensors compatibility: + - By default, all models are saved with the default pytorch serialization, so we use the list of default pytorch + files to know which safetensors files are needed. + - The model is safetensors compatible only if there is a matching safetensors file for every default pytorch file. + + Converting default pytorch serialized filenames to safetensors serialized filenames: + - For models from the diffusers library, just replace the ".bin" extension with ".safetensors" + - For models from the transformers library, the filename changes from "pytorch_model" to "model", and the ".bin" + extension is replaced with ".safetensors" + """ + pt_filenames = [] + + sf_filenames = set() + + passed_components = passed_components or [] + + for filename in filenames: + _, extension = os.path.splitext(filename) + + if len(filename.split("/")) == 2 and filename.split("/")[0] in passed_components: + continue + + if extension == ".bin": + pt_filenames.append(os.path.normpath(filename)) + elif extension == ".safetensors": + sf_filenames.add(os.path.normpath(filename)) + + for filename in pt_filenames: + # filename = 'foo/bar/baz.bam' -> path = 'foo/bar', filename = 'baz', extention = '.bam' + path, filename = os.path.split(filename) + filename, extension = os.path.splitext(filename) + + if filename.startswith("pytorch_model"): + filename = filename.replace("pytorch_model", "model") + else: + filename = filename + + expected_sf_filename = os.path.normpath(os.path.join(path, filename)) + expected_sf_filename = f"{expected_sf_filename}.safetensors" + if expected_sf_filename not in sf_filenames: + logger.warning(f"{expected_sf_filename} not found") + return False + + return True + + +def variant_compatible_siblings(filenames, variant=None) -> Union[List[os.PathLike], str]: + weight_names = [ + WEIGHTS_NAME, + SAFETENSORS_WEIGHTS_NAME, + FLAX_WEIGHTS_NAME, + ONNX_WEIGHTS_NAME, + ONNX_EXTERNAL_WEIGHTS_NAME, + ] + + if is_transformers_available(): + weight_names += [TRANSFORMERS_WEIGHTS_NAME, TRANSFORMERS_SAFE_WEIGHTS_NAME, TRANSFORMERS_FLAX_WEIGHTS_NAME] + + # model_pytorch, diffusion_model_pytorch, ... + weight_prefixes = [w.split(".")[0] for w in weight_names] + # .bin, .safetensors, ... + weight_suffixs = [w.split(".")[-1] for w in weight_names] + # -00001-of-00002 + transformers_index_format = r"\d{5}-of-\d{5}" + + if variant is not None: + # `diffusion_pytorch_model.fp16.bin` as well as `model.fp16-00001-of-00002.safetensors` + variant_file_re = re.compile( + rf"({'|'.join(weight_prefixes)})\.({variant}|{variant}-{transformers_index_format})\.({'|'.join(weight_suffixs)})$" + ) + # `text_encoder/pytorch_model.bin.index.fp16.json` + variant_index_re = re.compile( + rf"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.{variant}\.json$" + ) + + # `diffusion_pytorch_model.bin` as well as `model-00001-of-00002.safetensors` + non_variant_file_re = re.compile( + rf"({'|'.join(weight_prefixes)})(-{transformers_index_format})?\.({'|'.join(weight_suffixs)})$" + ) + # `text_encoder/pytorch_model.bin.index.json` + non_variant_index_re = re.compile(rf"({'|'.join(weight_prefixes)})\.({'|'.join(weight_suffixs)})\.index\.json") + + if variant is not None: + variant_weights = {f for f in filenames if variant_file_re.match(f.split("/")[-1]) is not None} + variant_indexes = {f for f in filenames if variant_index_re.match(f.split("/")[-1]) is not None} + variant_filenames = variant_weights | variant_indexes + else: + variant_filenames = set() + + non_variant_weights = {f for f in filenames if non_variant_file_re.match(f.split("/")[-1]) is not None} + non_variant_indexes = {f for f in filenames if non_variant_index_re.match(f.split("/")[-1]) is not None} + non_variant_filenames = non_variant_weights | non_variant_indexes + + # all variant filenames will be used by default + usable_filenames = set(variant_filenames) + + def convert_to_variant(filename): + if "index" in filename: + variant_filename = filename.replace("index", f"index.{variant}") + elif re.compile(f"^(.*?){transformers_index_format}").match(filename) is not None: + variant_filename = f"{filename.split('-')[0]}.{variant}-{'-'.join(filename.split('-')[1:])}" + else: + variant_filename = f"{filename.split('.')[0]}.{variant}.{filename.split('.')[1]}" + return variant_filename + + for f in non_variant_filenames: + variant_filename = convert_to_variant(f) + if variant_filename not in usable_filenames: + usable_filenames.add(f) + + return usable_filenames, variant_filenames + + +def warn_deprecated_model_variant(pretrained_model_name_or_path, use_auth_token, variant, revision, model_filenames): + info = model_info( + pretrained_model_name_or_path, + use_auth_token=use_auth_token, + revision=None, + ) + filenames = {sibling.rfilename for sibling in info.siblings} + comp_model_filenames, _ = variant_compatible_siblings(filenames, variant=revision) + comp_model_filenames = [".".join(f.split(".")[:1] + f.split(".")[2:]) for f in comp_model_filenames] + + if set(model_filenames).issubset(set(comp_model_filenames)): + warnings.warn( + f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'` even though you can load it via `variant=`{revision}`. Loading model variants via `revision='{revision}'` is deprecated and will be removed in diffusers v1. Please use `variant='{revision}'` instead.", + FutureWarning, + ) + else: + warnings.warn( + f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant='{revision}'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have the required variant filenames in the 'main' branch. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title '{pretrained_model_name_or_path} is missing {revision} files' so that the correct variant file can be added.", + FutureWarning, + ) + + +def _unwrap_model(model): + """Unwraps a model.""" + if is_compiled_module(model): + model = model._orig_mod + + if is_peft_available(): + from peft import PeftModel + + if isinstance(model, PeftModel): + model = model.base_model.model + + return model + + +def maybe_raise_or_warn( + library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module +): + """Simple helper method to raise or warn in case incorrect module has been passed""" + if not is_pipeline_module: + library = importlib.import_module(library_name) + class_obj = getattr(library, class_name) + class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} + + expected_class_obj = None + for class_name, class_candidate in class_candidates.items(): + if class_candidate is not None and issubclass(class_obj, class_candidate): + expected_class_obj = class_candidate + + # Dynamo wraps the original model in a private class. + # I didn't find a public API to get the original class. + sub_model = passed_class_obj[name] + unwrapped_sub_model = _unwrap_model(sub_model) + model_cls = unwrapped_sub_model.__class__ + + if not issubclass(model_cls, expected_class_obj): + raise ValueError( + f"{passed_class_obj[name]} is of type: {model_cls}, but should be" f" {expected_class_obj}" + ) + else: + logger.warning( + f"You have passed a non-standard module {passed_class_obj[name]}. We cannot verify whether it" + " has the correct type" + ) + + +def get_class_obj_and_candidates( + library_name, class_name, importable_classes, pipelines, is_pipeline_module, component_name=None, cache_dir=None +): + """Simple helper method to retrieve class object of module as well as potential parent class objects""" + component_folder = os.path.join(cache_dir, component_name) + + if is_pipeline_module: + pipeline_module = getattr(pipelines, library_name) + + class_obj = getattr(pipeline_module, class_name) + class_candidates = {c: class_obj for c in importable_classes.keys()} + elif os.path.isfile(os.path.join(component_folder, library_name + ".py")): + # load custom component + class_obj = get_class_from_dynamic_module( + component_folder, module_file=library_name + ".py", class_name=class_name + ) + class_candidates = {c: class_obj for c in importable_classes.keys()} + else: + if class_name == 'UNetSpatioTemporalConditionModel': + library_name = "ixformer.contrib.DeepCache.svd.unet_spatio_temporal_condition" + + # else we just import it from the library. + library = importlib.import_module(library_name) + + class_obj = getattr(library, class_name) + class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()} + + return class_obj, class_candidates + + +def _get_pipeline_class( + class_obj, + config, + load_connected_pipeline=False, + custom_pipeline=None, + repo_id=None, + hub_revision=None, + class_name=None, + cache_dir=None, + revision=None, +): + if custom_pipeline is not None: + if custom_pipeline.endswith(".py"): + path = Path(custom_pipeline) + # decompose into folder & file + file_name = path.name + custom_pipeline = path.parent.absolute() + elif repo_id is not None: + file_name = f"{custom_pipeline}.py" + custom_pipeline = repo_id + else: + file_name = CUSTOM_PIPELINE_FILE_NAME + + if repo_id is not None and hub_revision is not None: + # if we load the pipeline code from the Hub + # make sure to overwrite the `revison` + revision = hub_revision + + return get_class_from_dynamic_module( + custom_pipeline, + module_file=file_name, + class_name=class_name, + repo_id=repo_id, + cache_dir=cache_dir, + revision=revision, + ) + + if class_obj != DiffusionPipeline: + return class_obj + + diffusers_module = importlib.import_module(class_obj.__module__.split(".")[0]) + class_name = config["_class_name"] + class_name = class_name[4:] if class_name.startswith("Flax") else class_name + + pipeline_cls = getattr(diffusers_module, class_name) + + if load_connected_pipeline: + from .auto_pipeline import _get_connected_pipeline + + connected_pipeline_cls = _get_connected_pipeline(pipeline_cls) + if connected_pipeline_cls is not None: + logger.info( + f"Loading connected pipeline {connected_pipeline_cls.__name__} instead of {pipeline_cls.__name__} as specified via `load_connected_pipeline=True`" + ) + else: + logger.info(f"{pipeline_cls.__name__} has no connected pipeline class. Loading {pipeline_cls.__name__}.") + + pipeline_cls = connected_pipeline_cls or pipeline_cls + + return pipeline_cls + + +def load_sub_model( + library_name: str, + class_name: str, + importable_classes: List[Any], + pipelines: Any, + is_pipeline_module: bool, + pipeline_class: Any, + torch_dtype: torch.dtype, + provider: Any, + sess_options: Any, + device_map: Optional[Union[Dict[str, torch.device], str]], + max_memory: Optional[Dict[Union[int, str], Union[int, str]]], + offload_folder: Optional[Union[str, os.PathLike]], + offload_state_dict: bool, + model_variants: Dict[str, str], + name: str, + from_flax: bool, + variant: str, + low_cpu_mem_usage: bool, + cached_folder: Union[str, os.PathLike], + revision: str = None, +): + """Helper method to load the module `name` from `library_name` and `class_name`""" + # retrieve class candidates + class_obj, class_candidates = get_class_obj_and_candidates( + library_name, + class_name, + importable_classes, + pipelines, + is_pipeline_module, + component_name=name, + cache_dir=cached_folder, + ) + + load_method_name = None + # retrive load method name + for class_name, class_candidate in class_candidates.items(): + if class_candidate is not None and issubclass(class_obj, class_candidate): + load_method_name = importable_classes[class_name][1] + + # if load method name is None, then we have a dummy module -> raise Error + if load_method_name is None: + none_module = class_obj.__module__ + is_dummy_path = none_module.startswith(DUMMY_MODULES_FOLDER) or none_module.startswith( + TRANSFORMERS_DUMMY_MODULES_FOLDER + ) + if is_dummy_path and "dummy" in none_module: + # call class_obj for nice error message of missing requirements + class_obj() + + raise ValueError( + f"The component {class_obj} of {pipeline_class} cannot be loaded as it does not seem to have" + f" any of the loading methods defined in {ALL_IMPORTABLE_CLASSES}." + ) + + load_method = getattr(class_obj, load_method_name) + + # add kwargs to loading method + diffusers_module = importlib.import_module('diffusers')#__name__.split(".")[0]) + loading_kwargs = {} + if issubclass(class_obj, torch.nn.Module): + loading_kwargs["torch_dtype"] = torch_dtype + if issubclass(class_obj, diffusers_module.OnnxRuntimeModel): + loading_kwargs["provider"] = provider + loading_kwargs["sess_options"] = sess_options + + is_diffusers_model = issubclass(class_obj, diffusers_module.ModelMixin) + + if is_transformers_available(): + transformers_version = version.parse(version.parse(transformers.__version__).base_version) + else: + transformers_version = "N/A" + + is_transformers_model = ( + is_transformers_available() + and issubclass(class_obj, PreTrainedModel) + and transformers_version >= version.parse("4.20.0") + ) + + # When loading a transformers model, if the device_map is None, the weights will be initialized as opposed to diffusers. + # To make default loading faster we set the `low_cpu_mem_usage=low_cpu_mem_usage` flag which is `True` by default. + # This makes sure that the weights won't be initialized which significantly speeds up loading. + if is_diffusers_model or is_transformers_model: + loading_kwargs["device_map"] = device_map + loading_kwargs["max_memory"] = max_memory + loading_kwargs["offload_folder"] = offload_folder + loading_kwargs["offload_state_dict"] = offload_state_dict + loading_kwargs["variant"] = model_variants.pop(name, None) + if from_flax: + loading_kwargs["from_flax"] = True + + # the following can be deleted once the minimum required `transformers` version + # is higher than 4.27 + if ( + is_transformers_model + and loading_kwargs["variant"] is not None + and transformers_version < version.parse("4.27.0") + ): + raise ImportError( + f"When passing `variant='{variant}'`, please make sure to upgrade your `transformers` version to at least 4.27.0.dev0" + ) + elif is_transformers_model and loading_kwargs["variant"] is None: + loading_kwargs.pop("variant") + + # if `from_flax` and model is transformer model, can currently not load with `low_cpu_mem_usage` + if not (from_flax and is_transformers_model): + loading_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage + else: + loading_kwargs["low_cpu_mem_usage"] = False + + # check if the module is in a subdirectory + if os.path.isdir(os.path.join(cached_folder, name)): + loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs) + else: + # else load from the root directory + loaded_sub_model = load_method(cached_folder, **loading_kwargs) + + return loaded_sub_model + + +class DiffusionPipeline(ConfigMixin, PushToHubMixin): + r""" + Base class for all pipelines. + + [`DiffusionPipeline`] stores all components (models, schedulers, and processors) for diffusion pipelines and + provides methods for loading, downloading and saving models. It also includes methods to: + + - move all PyTorch modules to the device of your choice + - enable/disable the progress bar for the denoising iteration + + Class attributes: + + - **config_name** (`str`) -- The configuration filename that stores the class and module names of all the + diffusion pipeline's components. + - **_optional_components** (`List[str]`) -- List of all optional components that don't have to be passed to the + pipeline to function (should be overridden by subclasses). + """ + + config_name = "model_index.json" + model_cpu_offload_seq = None + _optional_components = [] + _exclude_from_cpu_offload = [] + _load_connected_pipes = False + _is_onnx = False + + def register_modules(self, **kwargs): + # import it here to avoid circular import + diffusers_module = importlib.import_module(__name__.split(".")[0]) + pipelines = getattr(diffusers_module, "svd") + + for name, module in kwargs.items(): + # retrieve library + if module is None or isinstance(module, (tuple, list)) and module[0] is None: + register_dict = {name: (None, None)} + else: + # register the config from the original module, not the dynamo compiled one + not_compiled_module = _unwrap_model(module) + + library = not_compiled_module.__module__.split(".")[0] + + # check if the module is a pipeline module + module_path_items = not_compiled_module.__module__.split(".") + pipeline_dir = module_path_items[-2] if len(module_path_items) > 2 else None + + path = not_compiled_module.__module__.split(".") + is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir) + + # if library is not in LOADABLE_CLASSES, then it is a custom module. + # Or if it's a pipeline module, then the module is inside the pipeline + # folder so we set the library to module name. + if is_pipeline_module: + library = pipeline_dir + elif library not in LOADABLE_CLASSES: + library = not_compiled_module.__module__ + + # retrieve class_name + class_name = not_compiled_module.__class__.__name__ + + register_dict = {name: (library, class_name)} + + # save model index config + self.register_to_config(**register_dict) + + # set models + setattr(self, name, module) + + def __setattr__(self, name: str, value: Any): + if name in self.__dict__ and hasattr(self.config, name): + # We need to overwrite the config if name exists in config + if isinstance(getattr(self.config, name), (tuple, list)): + if value is not None and self.config[name][0] is not None: + class_library_tuple = (value.__module__.split(".")[0], value.__class__.__name__) + else: + class_library_tuple = (None, None) + + self.register_to_config(**{name: class_library_tuple}) + else: + self.register_to_config(**{name: value}) + + super().__setattr__(name, value) + + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + safe_serialization: bool = True, + variant: Optional[str] = None, + push_to_hub: bool = False, + **kwargs, + ): + """ + Save all saveable variables of the pipeline to a directory. A pipeline variable can be saved and loaded if its + class implements both a save and loading method. The pipeline is easily reloaded using the + [`~DiffusionPipeline.from_pretrained`] class method. + + Arguments: + save_directory (`str` or `os.PathLike`): + Directory to save a pipeline to. Will be created if it doesn't exist. + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`. + variant (`str`, *optional*): + If specified, weights are saved in the format `pytorch_model..bin`. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + model_index_dict = dict(self.config) + model_index_dict.pop("_class_name", None) + model_index_dict.pop("_diffusers_version", None) + model_index_dict.pop("_module", None) + model_index_dict.pop("_name_or_path", None) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + private = kwargs.pop("private", False) + create_pr = kwargs.pop("create_pr", False) + token = kwargs.pop("token", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id + + expected_modules, optional_kwargs = self._get_signature_keys(self) + + def is_saveable_module(name, value): + if name not in expected_modules: + return False + if name in self._optional_components and value[0] is None: + return False + return True + + model_index_dict = {k: v for k, v in model_index_dict.items() if is_saveable_module(k, v)} + for pipeline_component_name in model_index_dict.keys(): + sub_model = getattr(self, pipeline_component_name) + model_cls = sub_model.__class__ + + # Dynamo wraps the original model in a private class. + # I didn't find a public API to get the original class. + if is_compiled_module(sub_model): + sub_model = _unwrap_model(sub_model) + model_cls = sub_model.__class__ + + save_method_name = None + # search for the model's base class in LOADABLE_CLASSES + for library_name, library_classes in LOADABLE_CLASSES.items(): + if library_name in sys.modules: + library = importlib.import_module(library_name) + else: + logger.info( + f"{library_name} is not installed. Cannot save {pipeline_component_name} as {library_classes} from {library_name}" + ) + + for base_class, save_load_methods in library_classes.items(): + class_candidate = getattr(library, base_class, None) + if class_candidate is not None and issubclass(model_cls, class_candidate): + # if we found a suitable base class in LOADABLE_CLASSES then grab its save method + save_method_name = save_load_methods[0] + break + if save_method_name is not None: + break + + if save_method_name is None: + logger.warn(f"self.{pipeline_component_name}={sub_model} of type {type(sub_model)} cannot be saved.") + # make sure that unsaveable components are not tried to be loaded afterward + self.register_to_config(**{pipeline_component_name: (None, None)}) + continue + + save_method = getattr(sub_model, save_method_name) + + # Call the save method with the argument safe_serialization only if it's supported + save_method_signature = inspect.signature(save_method) + save_method_accept_safe = "safe_serialization" in save_method_signature.parameters + save_method_accept_variant = "variant" in save_method_signature.parameters + + save_kwargs = {} + if save_method_accept_safe: + save_kwargs["safe_serialization"] = safe_serialization + if save_method_accept_variant: + save_kwargs["variant"] = variant + + save_method(os.path.join(save_directory, pipeline_component_name), **save_kwargs) + + # finally save the config + self.save_config(save_directory) + + if push_to_hub: + self._upload_folder( + save_directory, + repo_id, + token=token, + commit_message=commit_message, + create_pr=create_pr, + ) + + def to(self, *args, **kwargs): + r""" + Performs Pipeline dtype and/or device conversion. A torch.dtype and torch.device are inferred from the + arguments of `self.to(*args, **kwargs).` + + + + If the pipeline already has the correct torch.dtype and torch.device, then it is returned as is. Otherwise, + the returned pipeline is a copy of self with the desired torch.dtype and torch.device. + + + + + Here are the ways to call `to`: + + - `to(dtype, silence_dtype_warnings=False) → DiffusionPipeline` to return a pipeline with the specified + [`dtype`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype) + - `to(device, silence_dtype_warnings=False) → DiffusionPipeline` to return a pipeline with the specified + [`device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device) + - `to(device=None, dtype=None, silence_dtype_warnings=False) → DiffusionPipeline` to return a pipeline with the + specified [`device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device) and + [`dtype`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype) + + Arguments: + dtype (`torch.dtype`, *optional*): + Returns a pipeline with the specified + [`dtype`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype) + device (`torch.Device`, *optional*): + Returns a pipeline with the specified + [`device`](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device) + silence_dtype_warnings (`str`, *optional*, defaults to `False`): + Whether to omit warnings if the target `dtype` is not compatible with the target `device`. + + Returns: + [`DiffusionPipeline`]: The pipeline converted to specified `dtype` and/or `dtype`. + """ + + torch_dtype = kwargs.pop("torch_dtype", None) + if torch_dtype is not None: + deprecate("torch_dtype", "0.25.0", "") + torch_device = kwargs.pop("torch_device", None) + if torch_device is not None: + deprecate("torch_device", "0.25.0", "") + + dtype_kwarg = kwargs.pop("dtype", None) + device_kwarg = kwargs.pop("device", None) + silence_dtype_warnings = kwargs.pop("silence_dtype_warnings", False) + + if torch_dtype is not None and dtype_kwarg is not None: + raise ValueError( + "You have passed both `torch_dtype` and `dtype` as a keyword argument. Please make sure to only pass `dtype`." + ) + + dtype = torch_dtype or dtype_kwarg + + if torch_device is not None and device_kwarg is not None: + raise ValueError( + "You have passed both `torch_device` and `device` as a keyword argument. Please make sure to only pass `device`." + ) + + device = torch_device or device_kwarg + + dtype_arg = None + device_arg = None + if len(args) == 1: + if isinstance(args[0], torch.dtype): + dtype_arg = args[0] + else: + device_arg = torch.device(args[0]) if args[0] is not None else None + elif len(args) == 2: + if isinstance(args[0], torch.dtype): + raise ValueError( + "When passing two arguments, make sure the first corresponds to `device` and the second to `dtype`." + ) + device_arg = torch.device(args[0]) if args[0] is not None else None + dtype_arg = args[1] + elif len(args) > 2: + raise ValueError("Please make sure to pass at most two arguments (`device` and `dtype`) `.to(...)`") + + if dtype is not None and dtype_arg is not None: + raise ValueError( + "You have passed `dtype` both as an argument and as a keyword argument. Please only pass one of the two." + ) + + dtype = dtype or dtype_arg + + if device is not None and device_arg is not None: + raise ValueError( + "You have passed `device` both as an argument and as a keyword argument. Please only pass one of the two." + ) + + device = device or device_arg + + # throw warning if pipeline is in "offloaded"-mode but user tries to manually set to GPU. + def module_is_sequentially_offloaded(module): + if not is_accelerate_available() or is_accelerate_version("<", "0.14.0"): + return False + + return hasattr(module, "_hf_hook") and not isinstance( + module._hf_hook, (accelerate.hooks.CpuOffload, accelerate.hooks.AlignDevicesHook) + ) + + def module_is_offloaded(module): + if not is_accelerate_available() or is_accelerate_version("<", "0.17.0.dev0"): + return False + + return hasattr(module, "_hf_hook") and isinstance(module._hf_hook, accelerate.hooks.CpuOffload) + + # .to("cuda") would raise an error if the pipeline is sequentially offloaded, so we raise our own to make it clearer + pipeline_is_sequentially_offloaded = any( + module_is_sequentially_offloaded(module) for _, module in self.components.items() + ) + if pipeline_is_sequentially_offloaded and device and torch.device(device).type == "cuda": + raise ValueError( + "It seems like you have activated sequential model offloading by calling `enable_sequential_cpu_offload`, but are now attempting to move the pipeline to GPU. This is not compatible with offloading. Please, move your pipeline `.to('cpu')` or consider removing the move altogether if you use sequential offloading." + ) + + # Display a warning in this case (the operation succeeds but the benefits are lost) + pipeline_is_offloaded = any(module_is_offloaded(module) for _, module in self.components.items()) + if pipeline_is_offloaded and device and torch.device(device).type == "cuda": + logger.warning( + f"It seems like you have activated model offloading by calling `enable_model_cpu_offload`, but are now manually moving the pipeline to GPU. It is strongly recommended against doing so as memory gains from offloading are likely to be lost. Offloading automatically takes care of moving the individual components {', '.join(self.components.keys())} to GPU when needed. To make sure offloading works as expected, you should consider moving the pipeline back to CPU: `pipeline.to('cpu')` or removing the move altogether if you use offloading." + ) + + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + is_offloaded = pipeline_is_offloaded or pipeline_is_sequentially_offloaded + for module in modules: + is_loaded_in_8bit = hasattr(module, "is_loaded_in_8bit") and module.is_loaded_in_8bit + + if is_loaded_in_8bit and dtype is not None: + logger.warning( + f"The module '{module.__class__.__name__}' has been loaded in 8bit and conversion to {torch_dtype} is not yet supported. Module is still in 8bit precision." + ) + + if is_loaded_in_8bit and device is not None: + logger.warning( + f"The module '{module.__class__.__name__}' has been loaded in 8bit and moving it to {torch_dtype} via `.to()` is not yet supported. Module is still on {module.device}." + ) + else: + module.to(device, dtype) + + if ( + module.dtype == torch.float16 + and str(device) in ["cpu"] + and not silence_dtype_warnings + and not is_offloaded + ): + logger.warning( + "Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It" + " is not recommended to move them to `cpu` as running them will fail. Please make" + " sure to use an accelerator to run the pipeline in inference, due to the lack of" + " support for`float16` operations on this device in PyTorch. Please, remove the" + " `torch_dtype=torch.float16` argument, or use another device for inference." + ) + return self + + @property + def device(self) -> torch.device: + r""" + Returns: + `torch.device`: The torch device on which the pipeline is located. + """ + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + for module in modules: + return module.device + + return torch.device("cpu") + + @property + def dtype(self) -> torch.dtype: + r""" + Returns: + `torch.dtype`: The torch dtype on which the pipeline is located. + """ + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + for module in modules: + return module.dtype + + return torch.float32 + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): + r""" + Instantiate a PyTorch diffusion pipeline from pretrained pipeline weights. + + The pipeline is set in evaluation mode (`model.eval()`) by default. + + If you get the error message below, you need to finetune the weights for your downstream task: + + ``` + Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match: + - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated + You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. + ``` + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): + Can be either: + + - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline + hosted on the Hub. + - A path to a *directory* (for example `./my_pipeline_directory/`) containing pipeline weights + saved using + [`~DiffusionPipeline.save_pretrained`]. + torch_dtype (`str` or `torch.dtype`, *optional*): + Override the default `torch.dtype` and load the model with another dtype. If "auto" is passed, the + dtype is automatically derived from the model's weights. + custom_pipeline (`str`, *optional*): + + + + 🧪 This is an experimental feature and may change in the future. + + + + Can be either: + + - A string, the *repo id* (for example `hf-internal-testing/diffusers-dummy-pipeline`) of a custom + pipeline hosted on the Hub. The repository must contain a file called pipeline.py that defines + the custom pipeline. + - A string, the *file name* of a community pipeline hosted on GitHub under + [Community](https://github.com/huggingface/diffusers/tree/main/examples/community). Valid file + names must match the file name and not the pipeline script (`clip_guided_stable_diffusion` + instead of `clip_guided_stable_diffusion.py`). Community pipelines are always loaded from the + current main branch of GitHub. + - A path to a directory (`./my_pipeline_directory/`) containing a custom pipeline. The directory + must contain a file called `pipeline.py` that defines the custom pipeline. + + For more information on how to load and create custom pipelines, please have a look at [Loading and + Adding Custom + Pipelines](https://huggingface.co/docs/diffusers/using-diffusers/custom_pipeline_overview) + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory where a downloaded pretrained model configuration is cached if the standard cache + is not used. + resume_download (`bool`, *optional*, defaults to `False`): + Whether or not to resume downloading the model weights and configuration files. If set to `False`, any + incompletely downloaded files are deleted. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info(`bool`, *optional*, defaults to `False`): + Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files or not. If set to `True`, the model + won't be downloaded from the Hub. + use_auth_token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from + `diffusers-cli login` (stored in `~/.huggingface`) is used. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier + allowed by Git. + custom_revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id similar to + `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a + custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. + mirror (`str`, *optional*): + Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not + guarantee the timeliness or safety of the source, and you should refer to the mirror site for more + information. + device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*): + A map that specifies where each submodule should go. It doesn’t need to be defined for each + parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the + same device. + + Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For + more information about each option see [designing a device + map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). + max_memory (`Dict`, *optional*): + A dictionary device identifier for the maximum memory. Will default to the maximum memory available for + each GPU and the available CPU RAM if unset. + offload_folder (`str` or `os.PathLike`, *optional*): + The path to offload weights if device_map contains the value `"disk"`. + offload_state_dict (`bool`, *optional*): + If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if + the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True` + when there is some disk offload. + low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): + Speed up model loading only loading the pretrained weights and not initializing the weights. This also + tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. + Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this + argument to `True` will raise an error. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `None`, the safetensors weights are downloaded if they're available **and** if the + safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors + weights. If set to `False`, safetensors weights are not loaded. + use_onnx (`bool`, *optional*, defaults to `None`): + If set to `True`, ONNX weights will always be downloaded if present. If set to `False`, ONNX weights + will never be downloaded. By default `use_onnx` defaults to the `_is_onnx` class attribute which is + `False` for non-ONNX pipelines and `True` for ONNX pipelines. ONNX weights include both files ending + with `.onnx` and `.pb`. + kwargs (remaining dictionary of keyword arguments, *optional*): + Can be used to overwrite load and saveable variables (the pipeline components of the specific pipeline + class). The overwritten components are passed directly to the pipelines `__init__` method. See example + below for more information. + variant (`str`, *optional*): + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when + loading `from_flax`. + + + + To use private or [gated](https://huggingface.co/docs/hub/models-gated#gated-models) models, log-in with + `huggingface-cli login`. + + + + Examples: + + ```py + >>> from diffusers import DiffusionPipeline + + >>> # Download pipeline from huggingface.co and cache. + >>> pipeline = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256") + + >>> # Download pipeline that requires an authorization token + >>> # For more information on access tokens, please refer to this section + >>> # of the documentation](https://huggingface.co/docs/hub/security-tokens) + >>> pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5") + + >>> # Use a different scheduler + >>> from diffusers import LMSDiscreteScheduler + + >>> scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config) + >>> pipeline.scheduler = scheduler + ``` + """ + cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE) + resume_download = kwargs.pop("resume_download", False) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + from_flax = kwargs.pop("from_flax", False) + torch_dtype = kwargs.pop("torch_dtype", None) + custom_pipeline = kwargs.pop("custom_pipeline", None) + custom_revision = kwargs.pop("custom_revision", None) + provider = kwargs.pop("provider", None) + sess_options = kwargs.pop("sess_options", None) + device_map = kwargs.pop("device_map", None) + max_memory = kwargs.pop("max_memory", None) + offload_folder = kwargs.pop("offload_folder", None) + offload_state_dict = kwargs.pop("offload_state_dict", False) + low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None) + use_onnx = kwargs.pop("use_onnx", None) + load_connected_pipeline = kwargs.pop("load_connected_pipeline", False) + + # 1. Download the checkpoints and configs + # use snapshot download here to get it working from from_pretrained + if not os.path.isdir(pretrained_model_name_or_path): + if pretrained_model_name_or_path.count("/") > 1: + raise ValueError( + f'The provided pretrained_model_name_or_path "{pretrained_model_name_or_path}"' + " is neither a valid local path nor a valid repo id. Please check the parameter." + ) + cached_folder = cls.download( + pretrained_model_name_or_path, + cache_dir=cache_dir, + resume_download=resume_download, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + from_flax=from_flax, + use_safetensors=use_safetensors, + use_onnx=use_onnx, + custom_pipeline=custom_pipeline, + custom_revision=custom_revision, + variant=variant, + load_connected_pipeline=load_connected_pipeline, + **kwargs, + ) + else: + cached_folder = pretrained_model_name_or_path + + config_dict = cls.load_config(cached_folder) + + # pop out "_ignore_files" as it is only needed for download + config_dict.pop("_ignore_files", None) + + # 2. Define which model components should load variants + # We retrieve the information by matching whether variant + # model checkpoints exist in the subfolders + model_variants = {} + if variant is not None: + for folder in os.listdir(cached_folder): + folder_path = os.path.join(cached_folder, folder) + is_folder = os.path.isdir(folder_path) and folder in config_dict + variant_exists = is_folder and any( + p.split(".")[1].startswith(variant) for p in os.listdir(folder_path) + ) + if variant_exists: + model_variants[folder] = variant + + # 3. Load the pipeline class, if using custom module then load it from the hub + # if we load from explicit class, let's use it + custom_class_name = None + if os.path.isfile(os.path.join(cached_folder, f"{custom_pipeline}.py")): + custom_pipeline = os.path.join(cached_folder, f"{custom_pipeline}.py") + elif isinstance(config_dict["_class_name"], (list, tuple)) and os.path.isfile( + os.path.join(cached_folder, f"{config_dict['_class_name'][0]}.py") + ): + custom_pipeline = os.path.join(cached_folder, f"{config_dict['_class_name'][0]}.py") + custom_class_name = config_dict["_class_name"][1] + + pipeline_class = _get_pipeline_class( + cls, + config_dict, + load_connected_pipeline=load_connected_pipeline, + custom_pipeline=custom_pipeline, + class_name=custom_class_name, + cache_dir=cache_dir, + revision=custom_revision, + ) + + # DEPRECATED: To be removed in 1.0.0 + if pipeline_class.__name__ == "StableDiffusionInpaintPipeline" and version.parse( + version.parse(config_dict["_diffusers_version"]).base_version + ) <= version.parse("0.5.1"): + from diffusers import StableDiffusionInpaintPipeline, StableDiffusionInpaintPipelineLegacy + + pipeline_class = StableDiffusionInpaintPipelineLegacy + + deprecation_message = ( + "You are using a legacy checkpoint for inpainting with Stable Diffusion, therefore we are loading the" + f" {StableDiffusionInpaintPipelineLegacy} class instead of {StableDiffusionInpaintPipeline}. For" + " better inpainting results, we strongly suggest using Stable Diffusion's official inpainting" + " checkpoint: https://huggingface.co/runwayml/stable-diffusion-inpainting instead or adapting your" + f" checkpoint {pretrained_model_name_or_path} to the format of" + " https://huggingface.co/runwayml/stable-diffusion-inpainting. Note that we do not actively maintain" + " the {StableDiffusionInpaintPipelineLegacy} class and will likely remove it in version 1.0.0." + ) + deprecate("StableDiffusionInpaintPipelineLegacy", "1.0.0", deprecation_message, standard_warn=False) + + # 4. Define expected modules given pipeline signature + # and define non-None initialized modules (=`init_kwargs`) + + # some modules can be passed directly to the init + # in this case they are already instantiated in `kwargs` + # extract them here + expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class) + passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs} + passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs} + + init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs) + + # define init kwargs and make sure that optional component modules are filtered out + init_kwargs = { + k: init_dict.pop(k) + for k in optional_kwargs + if k in init_dict and k not in pipeline_class._optional_components + } + init_kwargs = {**init_kwargs, **passed_pipe_kwargs} + + # remove `null` components + def load_module(name, value): + if value[0] is None: + return False + if name in passed_class_obj and passed_class_obj[name] is None: + return False + return True + + init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)} + + # Special case: safety_checker must be loaded separately when using `from_flax` + if from_flax and "safety_checker" in init_dict and "safety_checker" not in passed_class_obj: + raise NotImplementedError( + "The safety checker cannot be automatically loaded when loading weights `from_flax`." + " Please, pass `safety_checker=None` to `from_pretrained`, and load the safety checker" + " separately if you need it." + ) + + # 5. Throw nice warnings / errors for fast accelerate loading + if len(unused_kwargs) > 0: + logger.warning( + f"Keyword arguments {unused_kwargs} are not expected by {pipeline_class.__name__} and will be ignored." + ) + + if low_cpu_mem_usage and not is_accelerate_available(): + low_cpu_mem_usage = False + logger.warning( + "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the" + " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install" + " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip" + " install accelerate\n```\n." + ) + + if device_map is not None and not is_torch_version(">=", "1.9.0"): + raise NotImplementedError( + "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set" + " `device_map=None`." + ) + + if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"): + raise NotImplementedError( + "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set" + " `low_cpu_mem_usage=False`." + ) + + if low_cpu_mem_usage is False and device_map is not None: + raise ValueError( + f"You cannot set `low_cpu_mem_usage` to False while using device_map={device_map} for loading and" + " dispatching. Please make sure to set `low_cpu_mem_usage=True`." + ) + + # import it here to avoid circular import + from diffusers import pipelines + + # 6. Load each module in the pipeline + for name, (library_name, class_name) in logging.tqdm(init_dict.items(), desc="Loading pipeline components..."): + # 6.1 - now that JAX/Flax is an official framework of the library, we might load from Flax names + class_name = class_name[4:] if class_name.startswith("Flax") else class_name + + # 6.2 Define all importable classes + is_pipeline_module = hasattr(pipelines, library_name) + importable_classes = ALL_IMPORTABLE_CLASSES + loaded_sub_model = None + + # 6.3 Use passed sub model or load class_name from library_name + if name in passed_class_obj: + # if the model is in a pipeline module, then we load it from the pipeline + # check that passed_class_obj has correct parent class + maybe_raise_or_warn( + library_name, library, class_name, importable_classes, passed_class_obj, name, is_pipeline_module + ) + + loaded_sub_model = passed_class_obj[name] + else: + # load sub model + loaded_sub_model = load_sub_model( + library_name=library_name, + class_name=class_name, + importable_classes=importable_classes, + pipelines=pipelines, + is_pipeline_module=is_pipeline_module, + pipeline_class=pipeline_class, + torch_dtype=torch_dtype, + provider=provider, + sess_options=sess_options, + device_map=device_map, + max_memory=max_memory, + offload_folder=offload_folder, + offload_state_dict=offload_state_dict, + model_variants=model_variants, + name=name, + from_flax=from_flax, + variant=variant, + low_cpu_mem_usage=low_cpu_mem_usage, + cached_folder=cached_folder, + revision=revision, + ) + logger.info( + f"Loaded {name} as {class_name} from `{name}` subfolder of {pretrained_model_name_or_path}." + ) + + init_kwargs[name] = loaded_sub_model # UNet(...), # DiffusionSchedule(...) + + if pipeline_class._load_connected_pipes and os.path.isfile(os.path.join(cached_folder, "README.md")): + modelcard = ModelCard.load(os.path.join(cached_folder, "README.md")) + connected_pipes = {prefix: getattr(modelcard.data, prefix, [None])[0] for prefix in CONNECTED_PIPES_KEYS} + load_kwargs = { + "cache_dir": cache_dir, + "resume_download": resume_download, + "force_download": force_download, + "proxies": proxies, + "local_files_only": local_files_only, + "use_auth_token": use_auth_token, + "revision": revision, + "torch_dtype": torch_dtype, + "custom_pipeline": custom_pipeline, + "custom_revision": custom_revision, + "provider": provider, + "sess_options": sess_options, + "device_map": device_map, + "max_memory": max_memory, + "offload_folder": offload_folder, + "offload_state_dict": offload_state_dict, + "low_cpu_mem_usage": low_cpu_mem_usage, + "variant": variant, + "use_safetensors": use_safetensors, + } + + def get_connected_passed_kwargs(prefix): + connected_passed_class_obj = { + k.replace(f"{prefix}_", ""): w for k, w in passed_class_obj.items() if k.split("_")[0] == prefix + } + connected_passed_pipe_kwargs = { + k.replace(f"{prefix}_", ""): w for k, w in passed_pipe_kwargs.items() if k.split("_")[0] == prefix + } + + connected_passed_kwargs = {**connected_passed_class_obj, **connected_passed_pipe_kwargs} + return connected_passed_kwargs + + connected_pipes = { + prefix: DiffusionPipeline.from_pretrained( + repo_id, **load_kwargs.copy(), **get_connected_passed_kwargs(prefix) + ) + for prefix, repo_id in connected_pipes.items() + if repo_id is not None + } + + for prefix, connected_pipe in connected_pipes.items(): + # add connected pipes to `init_kwargs` with _, e.g. "prior_text_encoder" + init_kwargs.update( + {"_".join([prefix, name]): component for name, component in connected_pipe.components.items()} + ) + + # 7. Potentially add passed objects if expected + missing_modules = set(expected_modules) - set(init_kwargs.keys()) + passed_modules = list(passed_class_obj.keys()) + optional_modules = pipeline_class._optional_components + if len(missing_modules) > 0 and missing_modules <= set(passed_modules + optional_modules): + for module in missing_modules: + init_kwargs[module] = passed_class_obj.get(module, None) + elif len(missing_modules) > 0: + passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys())) - optional_kwargs + raise ValueError( + f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed." + ) + + # 8. Instantiate the pipeline + model = pipeline_class(**init_kwargs) + + # 9. Save where the model was instantiated from + model.register_to_config(_name_or_path=pretrained_model_name_or_path) + return model + + @property + def name_or_path(self) -> str: + return getattr(self.config, "_name_or_path", None) + + @property + def _execution_device(self): + r""" + Returns the device on which the pipeline's models will be executed. After calling + [`~DiffusionPipeline.enable_sequential_cpu_offload`] the execution device can only be inferred from + Accelerate's module hooks. + """ + for name, model in self.components.items(): + if not isinstance(model, torch.nn.Module) or name in self._exclude_from_cpu_offload: + continue + + if not hasattr(model, "_hf_hook"): + return self.device + for module in model.modules(): + if ( + hasattr(module, "_hf_hook") + and hasattr(module._hf_hook, "execution_device") + and module._hf_hook.execution_device is not None + ): + return torch.device(module._hf_hook.execution_device) + return self.device + + def enable_model_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"): + r""" + Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared + to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward` + method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with + `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`. + + Arguments: + gpu_id (`int`, *optional*): + The ID of the accelerator that shall be used in inference. If not specified, it will default to 0. + device (`torch.Device` or `str`, *optional*, defaults to "cuda"): + The PyTorch device type of the accelerator that shall be used in inference. If not specified, it will + default to "cuda". + """ + if self.model_cpu_offload_seq is None: + raise ValueError( + "Model CPU offload cannot be enabled because no `model_cpu_offload_seq` class attribute is set." + ) + + if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"): + from accelerate import cpu_offload_with_hook + else: + raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.") + + torch_device = torch.device(device) + device_index = torch_device.index + + if gpu_id is not None and device_index is not None: + raise ValueError( + f"You have passed both `gpu_id`={gpu_id} and an index as part of the passed device `device`={device}" + f"Cannot pass both. Please make sure to either not define `gpu_id` or not pass the index as part of the device: `device`={torch_device.type}" + ) + + # _offload_gpu_id should be set to passed gpu_id (or id in passed `device`) or default to previously set id or default to 0 + self._offload_gpu_id = gpu_id or torch_device.index or getattr(self, "_offload_gpu_id", 0) + + device_type = torch_device.type + device = torch.device(f"{device_type}:{self._offload_gpu_id}") + + if self.device.type != "cpu": + self.to("cpu", silence_dtype_warnings=True) + device_mod = getattr(torch, self.device.type, None) + if hasattr(device_mod, "empty_cache") and device_mod.is_available(): + device_mod.empty_cache() # otherwise we don't see the memory savings (but they probably exist) + + all_model_components = {k: v for k, v in self.components.items() if isinstance(v, torch.nn.Module)} + + self._all_hooks = [] + hook = None + for model_str in self.model_cpu_offload_seq.split("->"): + model = all_model_components.pop(model_str, None) + if not isinstance(model, torch.nn.Module): + continue + + _, hook = cpu_offload_with_hook(model, device, prev_module_hook=hook) + self._all_hooks.append(hook) + + # CPU offload models that are not in the seq chain unless they are explicitly excluded + # these models will stay on CPU until maybe_free_model_hooks is called + # some models cannot be in the seq chain because they are iteratively called, such as controlnet + for name, model in all_model_components.items(): + if not isinstance(model, torch.nn.Module): + continue + + if name in self._exclude_from_cpu_offload: + model.to(device) + else: + _, hook = cpu_offload_with_hook(model, device) + self._all_hooks.append(hook) + + def maybe_free_model_hooks(self): + r""" + Function that offloads all components, removes all model hooks that were added when using + `enable_model_cpu_offload` and then applies them again. In case the model has not been offloaded this function + is a no-op. Make sure to add this function to the end of the `__call__` function of your pipeline so that it + functions correctly when applying enable_model_cpu_offload. + """ + if not hasattr(self, "_all_hooks") or len(self._all_hooks) == 0: + # `enable_model_cpu_offload` has not be called, so silently do nothing + return + + for hook in self._all_hooks: + # offload model and remove hook from model + hook.offload() + hook.remove() + + # make sure the model is in the same state as before calling it + self.enable_model_cpu_offload() + + def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda"): + r""" + Offloads all models to CPU using 🤗 Accelerate, significantly reducing memory usage. When called, the state + dicts of all `torch.nn.Module` components (except those in `self._exclude_from_cpu_offload`) are saved to CPU + and then moved to `torch.device('meta')` and loaded to GPU only when their specific submodule has its `forward` + method called. Offloading happens on a submodule basis. Memory savings are higher than with + `enable_model_cpu_offload`, but performance is lower. + + Arguments: + gpu_id (`int`, *optional*): + The ID of the accelerator that shall be used in inference. If not specified, it will default to 0. + device (`torch.Device` or `str`, *optional*, defaults to "cuda"): + The PyTorch device type of the accelerator that shall be used in inference. If not specified, it will + default to "cuda". + """ + if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"): + from accelerate import cpu_offload + else: + raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher") + + torch_device = torch.device(device) + device_index = torch_device.index + + if gpu_id is not None and device_index is not None: + raise ValueError( + f"You have passed both `gpu_id`={gpu_id} and an index as part of the passed device `device`={device}" + f"Cannot pass both. Please make sure to either not define `gpu_id` or not pass the index as part of the device: `device`={torch_device.type}" + ) + + # _offload_gpu_id should be set to passed gpu_id (or id in passed `device`) or default to previously set id or default to 0 + self._offload_gpu_id = gpu_id or torch_device.index or getattr(self, "_offload_gpu_id", 0) + + device_type = torch_device.type + device = torch.device(f"{device_type}:{self._offload_gpu_id}") + + if self.device.type != "cpu": + self.to("cpu", silence_dtype_warnings=True) + device_mod = getattr(torch, self.device.type, None) + if hasattr(device_mod, "empty_cache") and device_mod.is_available(): + device_mod.empty_cache() # otherwise we don't see the memory savings (but they probably exist) + + for name, model in self.components.items(): + if not isinstance(model, torch.nn.Module): + continue + + if name in self._exclude_from_cpu_offload: + model.to(device) + else: + # make sure to offload buffers if not all high level weights + # are of type nn.Module + offload_buffers = len(model._parameters) > 0 + cpu_offload(model, device, offload_buffers=offload_buffers) + + @classmethod + def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]: + r""" + Download and cache a PyTorch diffusion pipeline from pretrained pipeline weights. + + Parameters: + pretrained_model_name (`str` or `os.PathLike`, *optional*): + A string, the *repository id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline + hosted on the Hub. + custom_pipeline (`str`, *optional*): + Can be either: + + - A string, the *repository id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained + pipeline hosted on the Hub. The repository must contain a file called `pipeline.py` that defines + the custom pipeline. + + - A string, the *file name* of a community pipeline hosted on GitHub under + [Community](https://github.com/huggingface/diffusers/tree/main/examples/community). Valid file + names must match the file name and not the pipeline script (`clip_guided_stable_diffusion` + instead of `clip_guided_stable_diffusion.py`). Community pipelines are always loaded from the + current `main` branch of GitHub. + + - A path to a *directory* (`./my_pipeline_directory/`) containing a custom pipeline. The directory + must contain a file called `pipeline.py` that defines the custom pipeline. + + + + 🧪 This is an experimental feature and may change in the future. + + + + For more information on how to load and create custom pipelines, take a look at [How to contribute a + community pipeline](https://huggingface.co/docs/diffusers/main/en/using-diffusers/contribute_pipeline). + + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + resume_download (`bool`, *optional*, defaults to `False`): + Whether or not to resume downloading the model weights and configuration files. If set to `False`, any + incompletely downloaded files are deleted. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info(`bool`, *optional*, defaults to `False`): + Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only load local model weights and configuration files or not. If set to `True`, the model + won't be downloaded from the Hub. + use_auth_token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from + `diffusers-cli login` (stored in `~/.huggingface`) is used. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier + allowed by Git. + custom_revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id similar to + `revision` when loading a custom pipeline from the Hub. It can be a 🤗 Diffusers version when loading a + custom pipeline from GitHub, otherwise it defaults to `"main"` when loading from the Hub. + mirror (`str`, *optional*): + Mirror source to resolve accessibility issues if you're downloading a model in China. We do not + guarantee the timeliness or safety of the source, and you should refer to the mirror site for more + information. + variant (`str`, *optional*): + Load weights from a specified variant filename such as `"fp16"` or `"ema"`. This is ignored when + loading `from_flax`. + use_safetensors (`bool`, *optional*, defaults to `None`): + If set to `None`, the safetensors weights are downloaded if they're available **and** if the + safetensors library is installed. If set to `True`, the model is forcibly loaded from safetensors + weights. If set to `False`, safetensors weights are not loaded. + use_onnx (`bool`, *optional*, defaults to `False`): + If set to `True`, ONNX weights will always be downloaded if present. If set to `False`, ONNX weights + will never be downloaded. By default `use_onnx` defaults to the `_is_onnx` class attribute which is + `False` for non-ONNX pipelines and `True` for ONNX pipelines. ONNX weights include both files ending + with `.onnx` and `.pb`. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether or not to allow for custom pipelines and components defined on the Hub in their own files. This + option should only be set to `True` for repositories you trust and in which you have read the code, as + it will execute code present on the Hub on your local machine. + + Returns: + `os.PathLike`: + A path to the downloaded pipeline. + + + + To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with + `huggingface-cli login`. + + + + """ + cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE) + resume_download = kwargs.pop("resume_download", False) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + from_flax = kwargs.pop("from_flax", False) + custom_pipeline = kwargs.pop("custom_pipeline", None) + custom_revision = kwargs.pop("custom_revision", None) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None) + use_onnx = kwargs.pop("use_onnx", None) + load_connected_pipeline = kwargs.pop("load_connected_pipeline", False) + trust_remote_code = kwargs.pop("trust_remote_code", False) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = True + allow_pickle = True + + allow_patterns = None + ignore_patterns = None + + model_info_call_error: Optional[Exception] = None + if not local_files_only: + try: + info = model_info( + pretrained_model_name, + use_auth_token=use_auth_token, + revision=revision, + ) + except HTTPError as e: + logger.warn(f"Couldn't connect to the Hub: {e}.\nWill try to load from local cache.") + local_files_only = True + model_info_call_error = e # save error to reraise it if model is not cached locally + + if not local_files_only: + config_file = hf_hub_download( + pretrained_model_name, + cls.config_name, + cache_dir=cache_dir, + revision=revision, + proxies=proxies, + force_download=force_download, + resume_download=resume_download, + use_auth_token=use_auth_token, + ) + + config_dict = cls._dict_from_json_file(config_file) + ignore_filenames = config_dict.pop("_ignore_files", []) + + # retrieve all folder_names that contain relevant files + folder_names = [k for k, v in config_dict.items() if isinstance(v, list) and k != "_class_name"] + + filenames = {sibling.rfilename for sibling in info.siblings} + model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant) + diffusers_module = importlib.import_module(__name__.split(".")[0]) + #diffusers_module = importlib.import_module(cls.__module__.split(".")[0]) + pipelines = getattr(diffusers_module, "svd") + + # optionally create a custom component <> custom file mapping + custom_components = {} + for component in folder_names: + module_candidate = config_dict[component][0] + + if module_candidate is None or not isinstance(module_candidate, str): + continue + + # We compute candidate file path on the Hub. Do not use `os.path.join`. + candidate_file = f"{component}/{module_candidate}.py" + + if candidate_file in filenames: + custom_components[component] = module_candidate + elif module_candidate not in LOADABLE_CLASSES and not hasattr(pipelines, module_candidate): + raise ValueError( + f"{candidate_file} as defined in `model_index.json` does not exist in {pretrained_model_name} and is not a module in 'diffusers/pipelines'." + ) + + if len(variant_filenames) == 0 and variant is not None: + deprecation_message = ( + f"You are trying to load the model files of the `variant={variant}`, but no such modeling files are available." + f"The default model files: {model_filenames} will be loaded instead. Make sure to not load from `variant={variant}`" + "if such variant modeling files are not available. Doing so will lead to an error in v0.24.0 as defaulting to non-variant" + "modeling files is deprecated." + ) + deprecate("no variant default", "0.24.0", deprecation_message, standard_warn=False) + + # remove ignored filenames + model_filenames = set(model_filenames) - set(ignore_filenames) + variant_filenames = set(variant_filenames) - set(ignore_filenames) + + # if the whole pipeline is cached we don't have to ping the Hub + if revision in DEPRECATED_REVISION_ARGS and version.parse( + version.parse(__version__).base_version + ) >= version.parse("0.22.0"): + warn_deprecated_model_variant( + pretrained_model_name, use_auth_token, variant, revision, model_filenames + ) + + model_folder_names = {os.path.split(f)[0] for f in model_filenames if os.path.split(f)[0] in folder_names} + + custom_class_name = None + if custom_pipeline is None and isinstance(config_dict["_class_name"], (list, tuple)): + custom_pipeline = config_dict["_class_name"][0] + custom_class_name = config_dict["_class_name"][1] + + # all filenames compatible with variant will be added + allow_patterns = list(model_filenames) + + # allow all patterns from non-model folders + # this enables downloading schedulers, tokenizers, ... + allow_patterns += [f"{k}/*" for k in folder_names if k not in model_folder_names] + # add custom component files + allow_patterns += [f"{k}/{f}.py" for k, f in custom_components.items()] + # add custom pipeline file + allow_patterns += [f"{custom_pipeline}.py"] if f"{custom_pipeline}.py" in filenames else [] + # also allow downloading config.json files with the model + allow_patterns += [os.path.join(k, "config.json") for k in model_folder_names] + + allow_patterns += [ + SCHEDULER_CONFIG_NAME, + CONFIG_NAME, + cls.config_name, + CUSTOM_PIPELINE_FILE_NAME, + ] + + load_pipe_from_hub = custom_pipeline is not None and f"{custom_pipeline}.py" in filenames + load_components_from_hub = len(custom_components) > 0 + + if load_pipe_from_hub and not trust_remote_code: + raise ValueError( + f"The repository for {pretrained_model_name} contains custom code in {custom_pipeline}.py which must be executed to correctly " + f"load the model. You can inspect the repository content at https://hf.co/{pretrained_model_name}/blob/main/{custom_pipeline}.py.\n" + f"Please pass the argument `trust_remote_code=True` to allow custom code to be run." + ) + + if load_components_from_hub and not trust_remote_code: + raise ValueError( + f"The repository for {pretrained_model_name} contains custom code in {'.py, '.join([os.path.join(k, v) for k,v in custom_components.items()])} which must be executed to correctly " + f"load the model. You can inspect the repository content at {', '.join([f'https://hf.co/{pretrained_model_name}/{k}/{v}.py' for k,v in custom_components.items()])}.\n" + f"Please pass the argument `trust_remote_code=True` to allow custom code to be run." + ) + + # retrieve passed components that should not be downloaded + pipeline_class = _get_pipeline_class( + cls, + config_dict, + load_connected_pipeline=load_connected_pipeline, + custom_pipeline=custom_pipeline, + repo_id=pretrained_model_name if load_pipe_from_hub else None, + hub_revision=revision, + class_name=custom_class_name, + cache_dir=cache_dir, + revision=custom_revision, + ) + expected_components, _ = cls._get_signature_keys(pipeline_class) + passed_components = [k for k in expected_components if k in kwargs] + + if ( + use_safetensors + and not allow_pickle + and not is_safetensors_compatible( + model_filenames, variant=variant, passed_components=passed_components + ) + ): + raise EnvironmentError( + f"Could not find the necessary `safetensors` weights in {model_filenames} (variant={variant})" + ) + if from_flax: + ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb"] + elif use_safetensors and is_safetensors_compatible( + model_filenames, variant=variant, passed_components=passed_components + ): + ignore_patterns = ["*.bin", "*.msgpack"] + + use_onnx = use_onnx if use_onnx is not None else pipeline_class._is_onnx + if not use_onnx: + ignore_patterns += ["*.onnx", "*.pb"] + + safetensors_variant_filenames = {f for f in variant_filenames if f.endswith(".safetensors")} + safetensors_model_filenames = {f for f in model_filenames if f.endswith(".safetensors")} + if ( + len(safetensors_variant_filenames) > 0 + and safetensors_model_filenames != safetensors_variant_filenames + ): + logger.warn( + f"\nA mixture of {variant} and non-{variant} filenames will be loaded.\nLoaded {variant} filenames:\n[{', '.join(safetensors_variant_filenames)}]\nLoaded non-{variant} filenames:\n[{', '.join(safetensors_model_filenames - safetensors_variant_filenames)}\nIf this behavior is not expected, please check your folder structure." + ) + else: + ignore_patterns = ["*.safetensors", "*.msgpack"] + + use_onnx = use_onnx if use_onnx is not None else pipeline_class._is_onnx + if not use_onnx: + ignore_patterns += ["*.onnx", "*.pb"] + + bin_variant_filenames = {f for f in variant_filenames if f.endswith(".bin")} + bin_model_filenames = {f for f in model_filenames if f.endswith(".bin")} + if len(bin_variant_filenames) > 0 and bin_model_filenames != bin_variant_filenames: + logger.warn( + f"\nA mixture of {variant} and non-{variant} filenames will be loaded.\nLoaded {variant} filenames:\n[{', '.join(bin_variant_filenames)}]\nLoaded non-{variant} filenames:\n[{', '.join(bin_model_filenames - bin_variant_filenames)}\nIf this behavior is not expected, please check your folder structure." + ) + + # Don't download any objects that are passed + allow_patterns = [ + p for p in allow_patterns if not (len(p.split("/")) == 2 and p.split("/")[0] in passed_components) + ] + + if pipeline_class._load_connected_pipes: + allow_patterns.append("README.md") + + # Don't download index files of forbidden patterns either + ignore_patterns = ignore_patterns + [f"{i}.index.*json" for i in ignore_patterns] + + re_ignore_pattern = [re.compile(fnmatch.translate(p)) for p in ignore_patterns] + re_allow_pattern = [re.compile(fnmatch.translate(p)) for p in allow_patterns] + + expected_files = [f for f in filenames if not any(p.match(f) for p in re_ignore_pattern)] + expected_files = [f for f in expected_files if any(p.match(f) for p in re_allow_pattern)] + + snapshot_folder = Path(config_file).parent + pipeline_is_cached = all((snapshot_folder / f).is_file() for f in expected_files) + + if pipeline_is_cached and not force_download: + # if the pipeline is cached, we can directly return it + # else call snapshot_download + return snapshot_folder + + user_agent = {"pipeline_class": cls.__name__} + if custom_pipeline is not None and not custom_pipeline.endswith(".py"): + user_agent["custom_pipeline"] = custom_pipeline + + # download all allow_patterns - ignore_patterns + try: + cached_folder = snapshot_download( + pretrained_model_name, + cache_dir=cache_dir, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + user_agent=user_agent, + ) + + # retrieve pipeline class from local file + cls_name = cls.load_config(os.path.join(cached_folder, "model_index.json")).get("_class_name", None) + cls_name = cls_name[4:] if isinstance(cls_name, str) and cls_name.startswith("Flax") else cls_name + + diffusers_module = importlib.import_module(__name__.split(".")[0]) + pipeline_class = getattr(diffusers_module, cls_name, None) if isinstance(cls_name, str) else None + + if pipeline_class is not None and pipeline_class._load_connected_pipes: + modelcard = ModelCard.load(os.path.join(cached_folder, "README.md")) + connected_pipes = sum([getattr(modelcard.data, k, []) for k in CONNECTED_PIPES_KEYS], []) + for connected_pipe_repo_id in connected_pipes: + download_kwargs = { + "cache_dir": cache_dir, + "resume_download": resume_download, + "force_download": force_download, + "proxies": proxies, + "local_files_only": local_files_only, + "use_auth_token": use_auth_token, + "variant": variant, + "use_safetensors": use_safetensors, + } + DiffusionPipeline.download(connected_pipe_repo_id, **download_kwargs) + + return cached_folder + + except FileNotFoundError: + # Means we tried to load pipeline with `local_files_only=True` but the files have not been found in local cache. + # This can happen in two cases: + # 1. If the user passed `local_files_only=True` => we raise the error directly + # 2. If we forced `local_files_only=True` when `model_info` failed => we raise the initial error + if model_info_call_error is None: + # 1. user passed `local_files_only=True` + raise + else: + # 2. we forced `local_files_only=True` when `model_info` failed + raise EnvironmentError( + f"Cannot load model {pretrained_model_name}: model is not cached locally and an error occured" + " while trying to fetch metadata from the Hub. Please check out the root cause in the stacktrace" + " above." + ) from model_info_call_error + + @classmethod + def _get_signature_keys(cls, obj): + parameters = inspect.signature(obj.__init__).parameters + required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty} + optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty}) + expected_modules = set(required_parameters.keys()) - {"self"} + + optional_names = list(optional_parameters) + for name in optional_names: + if name in cls._optional_components: + expected_modules.add(name) + optional_parameters.remove(name) + + return expected_modules, optional_parameters + + @property + def components(self) -> Dict[str, Any]: + r""" + The `self.components` property can be useful to run different pipelines with the same weights and + configurations without reallocating additional memory. + + Returns (`dict`): + A dictionary containing all the modules needed to initialize the pipeline. + + Examples: + + ```py + >>> from diffusers import ( + ... StableDiffusionPipeline, + ... StableDiffusionImg2ImgPipeline, + ... StableDiffusionInpaintPipeline, + ... ) + + >>> text2img = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> img2img = StableDiffusionImg2ImgPipeline(**text2img.components) + >>> inpaint = StableDiffusionInpaintPipeline(**text2img.components) + ``` + """ + expected_modules, optional_parameters = self._get_signature_keys(self) + components = { + k: getattr(self, k) for k in self.config.keys() if not k.startswith("_") and k not in optional_parameters + } + + if set(components.keys()) != expected_modules: + raise ValueError( + f"{self} has been incorrectly initialized or {self.__class__} is incorrectly implemented. Expected" + f" {expected_modules} to be defined, but {components.keys()} are defined." + ) + + return components + + @staticmethod + def numpy_to_pil(images): + """ + Convert a NumPy image or a batch of images to a PIL image. + """ + return numpy_to_pil(images) + + def progress_bar(self, iterable=None, total=None): + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + elif not isinstance(self._progress_bar_config, dict): + raise ValueError( + f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." + ) + + if iterable is not None: + return tqdm(iterable, **self._progress_bar_config) + elif total is not None: + return tqdm(total=total, **self._progress_bar_config) + else: + raise ValueError("Either `total` or `iterable` has to be defined.") + + def set_progress_bar_config(self, **kwargs): + self._progress_bar_config = kwargs + + def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None): + r""" + Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). When this + option is enabled, you should observe lower GPU memory usage and a potential speed up during inference. Speed + up during training is not guaranteed. + + + + ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes + precedent. + + + + Parameters: + attention_op (`Callable`, *optional*): + Override the default `None` operator for use as `op` argument to the + [`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention) + function of xFormers. + + Examples: + + ```py + >>> import torch + >>> from diffusers import DiffusionPipeline + >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp + + >>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16) + >>> pipe = pipe.to("cuda") + >>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp) + >>> # Workaround for not accepting attention shape using VAE for Flash Attention + >>> pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None) + ``` + """ + self.set_use_memory_efficient_attention_xformers(True, attention_op) + + def disable_xformers_memory_efficient_attention(self): + r""" + Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). + """ + self.set_use_memory_efficient_attention_xformers(False) + + def set_use_memory_efficient_attention_xformers( + self, valid: bool, attention_op: Optional[Callable] = None + ) -> None: + # Recursively walk through all the children. + # Any children which exposes the set_use_memory_efficient_attention_xformers method + # gets the message + def fn_recursive_set_mem_eff(module: torch.nn.Module): + if hasattr(module, "set_use_memory_efficient_attention_xformers"): + module.set_use_memory_efficient_attention_xformers(valid, attention_op) + + for child in module.children(): + fn_recursive_set_mem_eff(child) + + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module)] + + for module in modules: + fn_recursive_set_mem_eff(module) + + def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"): + r""" + Enable sliced attention computation. When this option is enabled, the attention module splits the input tensor + in slices to compute attention in several steps. For more than one attention head, the computation is performed + sequentially over each head. This is useful to save some memory in exchange for a small speed decrease. + + + + ⚠️ Don't enable attention slicing if you're already using `scaled_dot_product_attention` (SDPA) from PyTorch + 2.0 or xFormers. These attention computations are already very memory efficient so you won't need to enable + this function. If you enable attention slicing with SDPA or xFormers, it can lead to serious slow downs! + + + + Args: + slice_size (`str` or `int`, *optional*, defaults to `"auto"`): + When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If + `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + + Examples: + + ```py + >>> import torch + >>> from diffusers import StableDiffusionPipeline + + >>> pipe = StableDiffusionPipeline.from_pretrained( + ... "runwayml/stable-diffusion-v1-5", + ... torch_dtype=torch.float16, + ... use_safetensors=True, + ... ) + + >>> prompt = "a photo of an astronaut riding a horse on mars" + >>> pipe.enable_attention_slicing() + >>> image = pipe(prompt).images[0] + ``` + """ + self.set_attention_slice(slice_size) + + def disable_attention_slicing(self): + r""" + Disable sliced attention computation. If `enable_attention_slicing` was previously called, attention is + computed in one step. + """ + # set slice_size = `None` to disable `attention slicing` + self.enable_attention_slicing(None) + + def set_attention_slice(self, slice_size: Optional[int]): + module_names, _ = self._get_signature_keys(self) + modules = [getattr(self, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module) and hasattr(m, "set_attention_slice")] + + for module in modules: + module.set_attention_slice(slice_size) diff --git a/ixformer_sdk/contrib/DeepCache/svd/unet_3d_blocks.py b/ixformer_sdk/contrib/DeepCache/svd/unet_3d_blocks.py new file mode 100644 index 00000000..fe0ab0e9 --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/svd/unet_3d_blocks.py @@ -0,0 +1,2412 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. + +from typing import Any, Dict, Optional, Tuple, Union + +import torch +from torch import nn + +from diffusers.utils import is_torch_version +from diffusers.utils.torch_utils import apply_freeu +from diffusers.models.attention import Attention +from diffusers.models.dual_transformer_2d import DualTransformer2DModel +from diffusers.models.resnet import ( + Downsample2D, + ResnetBlock2D, + SpatioTemporalResBlock, + TemporalConvLayer, + Upsample2D, +) +from diffusers.models.transformer_2d import Transformer2DModel +from diffusers.models.transformer_temporal import ( + TransformerSpatioTemporalModel, + TransformerTemporalModel, +) + + +def get_down_block( + down_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + temb_channels: int, + add_downsample: bool, + resnet_eps: float, + resnet_act_fn: str, + num_attention_heads: int, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + downsample_padding: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = True, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + temporal_num_attention_heads: int = 8, + temporal_max_seq_length: int = 32, + transformer_layers_per_block: int = 1, +) -> Union[ + "DownBlock3D", + "CrossAttnDownBlock3D", + "DownBlockMotion", + "CrossAttnDownBlockMotion", + "DownBlockSpatioTemporal", + "CrossAttnDownBlockSpatioTemporal", +]: + if down_block_type == "DownBlock3D": + return DownBlock3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "CrossAttnDownBlock3D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock3D") + return CrossAttnDownBlock3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + if down_block_type == "DownBlockMotion": + return DownBlockMotion( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + temporal_num_attention_heads=temporal_num_attention_heads, + temporal_max_seq_length=temporal_max_seq_length, + ) + elif down_block_type == "CrossAttnDownBlockMotion": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlockMotion") + return CrossAttnDownBlockMotion( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + temporal_num_attention_heads=temporal_num_attention_heads, + temporal_max_seq_length=temporal_max_seq_length, + ) + elif down_block_type == "DownBlockSpatioTemporal": + # added for SDV + return DownBlockSpatioTemporal( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + ) + elif down_block_type == "CrossAttnDownBlockSpatioTemporal": + # added for SDV + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlockSpatioTemporal") + return CrossAttnDownBlockSpatioTemporal( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + num_layers=num_layers, + transformer_layers_per_block=transformer_layers_per_block, + add_downsample=add_downsample, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + ) + + raise ValueError(f"{down_block_type} does not exist.") + + +def get_up_block( + up_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + add_upsample: bool, + resnet_eps: float, + resnet_act_fn: str, + num_attention_heads: int, + resolution_idx: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = True, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + temporal_num_attention_heads: int = 8, + temporal_cross_attention_dim: Optional[int] = None, + temporal_max_seq_length: int = 32, + transformer_layers_per_block: int = 1, + dropout: float = 0.0, +) -> Union[ + "UpBlock3D", + "CrossAttnUpBlock3D", + "UpBlockMotion", + "CrossAttnUpBlockMotion", + "UpBlockSpatioTemporal", + "CrossAttnUpBlockSpatioTemporal", +]: + if up_block_type == "UpBlock3D": + return UpBlock3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + resolution_idx=resolution_idx, + ) + elif up_block_type == "CrossAttnUpBlock3D": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock3D") + return CrossAttnUpBlock3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + resolution_idx=resolution_idx, + ) + if up_block_type == "UpBlockMotion": + return UpBlockMotion( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + resolution_idx=resolution_idx, + temporal_num_attention_heads=temporal_num_attention_heads, + temporal_max_seq_length=temporal_max_seq_length, + ) + elif up_block_type == "CrossAttnUpBlockMotion": + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlockMotion") + return CrossAttnUpBlockMotion( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + resolution_idx=resolution_idx, + temporal_num_attention_heads=temporal_num_attention_heads, + temporal_max_seq_length=temporal_max_seq_length, + ) + elif up_block_type == "UpBlockSpatioTemporal": + # added for SDV + return UpBlockSpatioTemporal( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + resolution_idx=resolution_idx, + add_upsample=add_upsample, + ) + elif up_block_type == "CrossAttnUpBlockSpatioTemporal": + # added for SDV + if cross_attention_dim is None: + raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlockSpatioTemporal") + return CrossAttnUpBlockSpatioTemporal( + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + num_layers=num_layers, + transformer_layers_per_block=transformer_layers_per_block, + add_upsample=add_upsample, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + resolution_idx=resolution_idx, + ) + + raise ValueError(f"{up_block_type} does not exist.") + + +class UNetMidBlock3DCrossAttn(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads: int = 1, + output_scale_factor: float = 1.0, + cross_attention_dim: int = 1280, + dual_cross_attention: bool = False, + use_linear_projection: bool = True, + upcast_attention: bool = False, + ): + super().__init__() + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + temp_convs = [ + TemporalConvLayer( + in_channels, + in_channels, + dropout=0.1, + norm_num_groups=resnet_groups, + ) + ] + attentions = [] + temp_attentions = [] + + for _ in range(num_layers): + attentions.append( + Transformer2DModel( + in_channels // num_attention_heads, + num_attention_heads, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + ) + ) + temp_attentions.append( + TransformerTemporalModel( + in_channels // num_attention_heads, + num_attention_heads, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + temp_convs.append( + TemporalConvLayer( + in_channels, + in_channels, + dropout=0.1, + norm_num_groups=resnet_groups, + ) + ) + + self.resnets = nn.ModuleList(resnets) + self.temp_convs = nn.ModuleList(temp_convs) + self.attentions = nn.ModuleList(attentions) + self.temp_attentions = nn.ModuleList(temp_attentions) + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + num_frames: int = 1, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + ) -> torch.FloatTensor: + hidden_states = self.resnets[0](hidden_states, temb) + hidden_states = self.temp_convs[0](hidden_states, num_frames=num_frames) + for attn, temp_attn, resnet, temp_conv in zip( + self.attentions, self.temp_attentions, self.resnets[1:], self.temp_convs[1:] + ): + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + hidden_states = temp_attn( + hidden_states, + num_frames=num_frames, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + hidden_states = resnet(hidden_states, temb) + hidden_states = temp_conv(hidden_states, num_frames=num_frames) + + return hidden_states + + +class CrossAttnDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads: int = 1, + cross_attention_dim: int = 1280, + output_scale_factor: float = 1.0, + downsample_padding: int = 1, + add_downsample: bool = True, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + ): + super().__init__() + resnets = [] + attentions = [] + temp_attentions = [] + temp_convs = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + temp_convs.append( + TemporalConvLayer( + out_channels, + out_channels, + dropout=0.1, + norm_num_groups=resnet_groups, + ) + ) + attentions.append( + Transformer2DModel( + out_channels // num_attention_heads, + num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + ) + ) + temp_attentions.append( + TransformerTemporalModel( + out_channels // num_attention_heads, + num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + self.resnets = nn.ModuleList(resnets) + self.temp_convs = nn.ModuleList(temp_convs) + self.attentions = nn.ModuleList(attentions) + self.temp_attentions = nn.ModuleList(temp_attentions) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + num_frames: int = 1, + cross_attention_kwargs: Dict[str, Any] = None, + ) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: + # TODO(Patrick, William) - attention mask is not used + output_states = () + + for resnet, temp_conv, attn, temp_attn in zip( + self.resnets, self.temp_convs, self.attentions, self.temp_attentions + ): + hidden_states = resnet(hidden_states, temb) + hidden_states = temp_conv(hidden_states, num_frames=num_frames) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + hidden_states = temp_attn( + hidden_states, + num_frames=num_frames, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states += (hidden_states,) + + return hidden_states, output_states + + +class DownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + downsample_padding: int = 1, + ): + super().__init__() + resnets = [] + temp_convs = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + temp_convs.append( + TemporalConvLayer( + out_channels, + out_channels, + dropout=0.1, + norm_num_groups=resnet_groups, + ) + ) + + self.resnets = nn.ModuleList(resnets) + self.temp_convs = nn.ModuleList(temp_convs) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + num_frames: int = 1, + ) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: + output_states = () + + for resnet, temp_conv in zip(self.resnets, self.temp_convs): + hidden_states = resnet(hidden_states, temb) + hidden_states = temp_conv(hidden_states, num_frames=num_frames) + + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states += (hidden_states,) + + return hidden_states, output_states + + +class CrossAttnUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads: int = 1, + cross_attention_dim: int = 1280, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resolution_idx: Optional[int] = None, + ): + super().__init__() + resnets = [] + temp_convs = [] + attentions = [] + temp_attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + temp_convs.append( + TemporalConvLayer( + out_channels, + out_channels, + dropout=0.1, + norm_num_groups=resnet_groups, + ) + ) + attentions.append( + Transformer2DModel( + out_channels // num_attention_heads, + num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + ) + ) + temp_attentions.append( + TransformerTemporalModel( + out_channels // num_attention_heads, + num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + self.resnets = nn.ModuleList(resnets) + self.temp_convs = nn.ModuleList(temp_convs) + self.attentions = nn.ModuleList(attentions) + self.temp_attentions = nn.ModuleList(temp_attentions) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + num_frames: int = 1, + cross_attention_kwargs: Dict[str, Any] = None, + ) -> torch.FloatTensor: + is_freeu_enabled = ( + getattr(self, "s1", None) + and getattr(self, "s2", None) + and getattr(self, "b1", None) + and getattr(self, "b2", None) + ) + + # TODO(Patrick, William) - attention mask is not used + for resnet, temp_conv, attn, temp_attn in zip( + self.resnets, self.temp_convs, self.attentions, self.temp_attentions + ): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # FreeU: Only operate on the first two stages + if is_freeu_enabled: + hidden_states, res_hidden_states = apply_freeu( + self.resolution_idx, + hidden_states, + res_hidden_states, + s1=self.s1, + s2=self.s2, + b1=self.b1, + b2=self.b2, + ) + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb) + hidden_states = temp_conv(hidden_states, num_frames=num_frames) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + hidden_states = temp_attn( + hidden_states, + num_frames=num_frames, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + +class UpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + resolution_idx: Optional[int] = None, + ): + super().__init__() + resnets = [] + temp_convs = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + temp_convs.append( + TemporalConvLayer( + out_channels, + out_channels, + dropout=0.1, + norm_num_groups=resnet_groups, + ) + ) + + self.resnets = nn.ModuleList(resnets) + self.temp_convs = nn.ModuleList(temp_convs) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + upsample_size: Optional[int] = None, + num_frames: int = 1, + ) -> torch.FloatTensor: + is_freeu_enabled = ( + getattr(self, "s1", None) + and getattr(self, "s2", None) + and getattr(self, "b1", None) + and getattr(self, "b2", None) + ) + for resnet, temp_conv in zip(self.resnets, self.temp_convs): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # FreeU: Only operate on the first two stages + if is_freeu_enabled: + hidden_states, res_hidden_states = apply_freeu( + self.resolution_idx, + hidden_states, + res_hidden_states, + s1=self.s1, + s2=self.s2, + b1=self.b1, + b2=self.b2, + ) + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + hidden_states = resnet(hidden_states, temb) + hidden_states = temp_conv(hidden_states, num_frames=num_frames) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + +class DownBlockMotion(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + downsample_padding: int = 1, + temporal_num_attention_heads: int = 1, + temporal_cross_attention_dim: Optional[int] = None, + temporal_max_seq_length: int = 32, + ): + super().__init__() + resnets = [] + motion_modules = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + motion_modules.append( + TransformerTemporalModel( + num_attention_heads=temporal_num_attention_heads, + in_channels=out_channels, + norm_num_groups=resnet_groups, + cross_attention_dim=temporal_cross_attention_dim, + attention_bias=False, + activation_fn="geglu", + positional_embeddings="sinusoidal", + num_positional_embeddings=temporal_max_seq_length, + attention_head_dim=out_channels // temporal_num_attention_heads, + ) + ) + + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + num_frames: int = 1, + ) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: + output_states = () + + blocks = zip(self.resnets, self.motion_modules) + for resnet, motion_module in blocks: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + use_reentrant=False, + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb, scale + ) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(motion_module), + hidden_states.requires_grad_(), + temb, + num_frames, + ) + + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + hidden_states = motion_module(hidden_states, num_frames=num_frames)[0] + + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale=scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class CrossAttnDownBlockMotion(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads: int = 1, + cross_attention_dim: int = 1280, + output_scale_factor: float = 1.0, + downsample_padding: int = 1, + add_downsample: bool = True, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + attention_type: str = "default", + temporal_cross_attention_dim: Optional[int] = None, + temporal_num_attention_heads: int = 8, + temporal_max_seq_length: int = 32, + ): + super().__init__() + resnets = [] + attentions = [] + motion_modules = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + + motion_modules.append( + TransformerTemporalModel( + num_attention_heads=temporal_num_attention_heads, + in_channels=out_channels, + norm_num_groups=resnet_groups, + cross_attention_dim=temporal_cross_attention_dim, + attention_bias=False, + activation_fn="geglu", + positional_embeddings="sinusoidal", + num_positional_embeddings=temporal_max_seq_length, + attention_head_dim=out_channels // temporal_num_attention_heads, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + num_frames: int = 1, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + additional_residuals: Optional[torch.FloatTensor] = None, + ): + output_states = () + + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + blocks = list(zip(self.resnets, self.attentions, self.motion_modules)) + for i, (resnet, attn, motion_module) in enumerate(blocks): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + hidden_states = motion_module( + hidden_states, + num_frames=num_frames, + )[0] + + # apply additional residuals to the output of the last pair of resnet and attention blocks + if i == len(blocks) - 1 and additional_residuals is not None: + hidden_states = hidden_states + additional_residuals + + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale=lora_scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class CrossAttnUpBlockMotion(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + resolution_idx: Optional[int] = None, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads: int = 1, + cross_attention_dim: int = 1280, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + attention_type: str = "default", + temporal_cross_attention_dim: Optional[int] = None, + temporal_num_attention_heads: int = 8, + temporal_max_seq_length: int = 32, + ): + super().__init__() + resnets = [] + attentions = [] + motion_modules = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + motion_modules.append( + TransformerTemporalModel( + num_attention_heads=temporal_num_attention_heads, + in_channels=out_channels, + norm_num_groups=resnet_groups, + cross_attention_dim=temporal_cross_attention_dim, + attention_bias=False, + activation_fn="geglu", + positional_embeddings="sinusoidal", + num_positional_embeddings=temporal_max_seq_length, + attention_head_dim=out_channels // temporal_num_attention_heads, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + num_frames: int = 1, + ) -> torch.FloatTensor: + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + is_freeu_enabled = ( + getattr(self, "s1", None) + and getattr(self, "s2", None) + and getattr(self, "b1", None) + and getattr(self, "b2", None) + ) + + blocks = zip(self.resnets, self.attentions, self.motion_modules) + for resnet, attn, motion_module in blocks: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # FreeU: Only operate on the first two stages + if is_freeu_enabled: + hidden_states, res_hidden_states = apply_freeu( + self.resolution_idx, + hidden_states, + res_hidden_states, + s1=self.s1, + s2=self.s2, + b1=self.b1, + b2=self.b2, + ) + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + hidden_states = motion_module( + hidden_states, + num_frames=num_frames, + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size, scale=lora_scale) + + return hidden_states + + +class UpBlockMotion(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + resolution_idx: Optional[int] = None, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + temporal_norm_num_groups: int = 32, + temporal_cross_attention_dim: Optional[int] = None, + temporal_num_attention_heads: int = 8, + temporal_max_seq_length: int = 32, + ): + super().__init__() + resnets = [] + motion_modules = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + motion_modules.append( + TransformerTemporalModel( + num_attention_heads=temporal_num_attention_heads, + in_channels=out_channels, + norm_num_groups=temporal_norm_num_groups, + cross_attention_dim=temporal_cross_attention_dim, + attention_bias=False, + activation_fn="geglu", + positional_embeddings="sinusoidal", + num_positional_embeddings=temporal_max_seq_length, + attention_head_dim=out_channels // temporal_num_attention_heads, + ) + ) + + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + upsample_size=None, + scale: float = 1.0, + num_frames: int = 1, + ) -> torch.FloatTensor: + is_freeu_enabled = ( + getattr(self, "s1", None) + and getattr(self, "s2", None) + and getattr(self, "b1", None) + and getattr(self, "b2", None) + ) + + blocks = zip(self.resnets, self.motion_modules) + + for resnet, motion_module in blocks: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # FreeU: Only operate on the first two stages + if is_freeu_enabled: + hidden_states, res_hidden_states = apply_freeu( + self.resolution_idx, + hidden_states, + res_hidden_states, + s1=self.s1, + s2=self.s2, + b1=self.b1, + b2=self.b2, + ) + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + use_reentrant=False, + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + ) + + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + hidden_states = motion_module(hidden_states, num_frames=num_frames)[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size, scale=scale) + + return hidden_states + + +class UNetMidBlockCrossAttnMotion(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads: int = 1, + output_scale_factor: float = 1.0, + cross_attention_dim: int = 1280, + dual_cross_attention: float = False, + use_linear_projection: float = False, + upcast_attention: float = False, + attention_type: str = "default", + temporal_num_attention_heads: int = 1, + temporal_cross_attention_dim: Optional[int] = None, + temporal_max_seq_length: int = 32, + ): + super().__init__() + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + motion_modules = [] + + for _ in range(num_layers): + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + in_channels // num_attention_heads, + in_channels=in_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + in_channels // num_attention_heads, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + motion_modules.append( + TransformerTemporalModel( + num_attention_heads=temporal_num_attention_heads, + attention_head_dim=in_channels // temporal_num_attention_heads, + in_channels=in_channels, + norm_num_groups=resnet_groups, + cross_attention_dim=temporal_cross_attention_dim, + attention_bias=False, + positional_embeddings="sinusoidal", + num_positional_embeddings=temporal_max_seq_length, + activation_fn="geglu", + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + num_frames: int = 1, + ) -> torch.FloatTensor: + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale) + + blocks = zip(self.attentions, self.resnets[1:], self.motion_modules) + for attn, resnet, motion_module in blocks: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(motion_module), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + else: + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + hidden_states = motion_module( + hidden_states, + num_frames=num_frames, + )[0] + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + return hidden_states + + +class MidBlockTemporalDecoder(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + attention_head_dim: int = 512, + num_layers: int = 1, + upcast_attention: bool = False, + ): + super().__init__() + + resnets = [] + attentions = [] + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + resnets.append( + SpatioTemporalResBlock( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=None, + eps=1e-6, + temporal_eps=1e-5, + merge_factor=0.0, + merge_strategy="learned", + switch_spatial_to_temporal_mix=True, + ) + ) + + attentions.append( + Attention( + query_dim=in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + eps=1e-6, + upcast_attention=upcast_attention, + norm_num_groups=32, + bias=True, + residual_connection=True, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward( + self, + hidden_states: torch.FloatTensor, + image_only_indicator: torch.FloatTensor, + ): + hidden_states = self.resnets[0]( + hidden_states, + image_only_indicator=image_only_indicator, + ) + for resnet, attn in zip(self.resnets[1:], self.attentions): + hidden_states = attn(hidden_states) + hidden_states = resnet( + hidden_states, + image_only_indicator=image_only_indicator, + ) + + return hidden_states + + +class UpBlockTemporalDecoder(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + add_upsample: bool = True, + ): + super().__init__() + resnets = [] + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + SpatioTemporalResBlock( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=None, + eps=1e-6, + temporal_eps=1e-5, + merge_factor=0.0, + merge_strategy="learned", + switch_spatial_to_temporal_mix=True, + ) + ) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + def forward( + self, + hidden_states: torch.FloatTensor, + image_only_indicator: torch.FloatTensor, + ) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet( + hidden_states, + image_only_indicator=image_only_indicator, + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class UNetMidBlockSpatioTemporal(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + num_layers: int = 1, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + num_attention_heads: int = 1, + cross_attention_dim: int = 1280, + ): + super().__init__() + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + # support for variable transformer layers per block + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * num_layers + + # there is always at least one resnet + resnets = [ + SpatioTemporalResBlock( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=1e-5, + ) + ] + attentions = [] + + for i in range(num_layers): + attentions.append( + TransformerSpatioTemporalModel( + num_attention_heads, + in_channels // num_attention_heads, + in_channels=in_channels, + num_layers=transformer_layers_per_block[i], + cross_attention_dim=cross_attention_dim, + ) + ) + + resnets.append( + SpatioTemporalResBlock( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=1e-5, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + image_only_indicator: Optional[torch.Tensor] = None, + ) -> torch.FloatTensor: + hidden_states = self.resnets[0]( + hidden_states, + temb, + image_only_indicator=image_only_indicator, + ) + + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if self.training and self.gradient_checkpointing: # TODO + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + return_dict=False, + )[0] + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + **ckpt_kwargs, + ) + else: + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + return_dict=False, + )[0] + hidden_states = resnet( + hidden_states, + temb, + image_only_indicator=image_only_indicator, + ) + + return hidden_states + + +class DownBlockSpatioTemporal(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + num_layers: int = 1, + add_downsample: bool = True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + SpatioTemporalResBlock( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=1e-5, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + image_only_indicator: Optional[torch.Tensor] = None, + exist_module_idx: Optional[int] = None, + ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: + output_states = () + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + use_reentrant=False, + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + ) + else: + hidden_states = resnet( + hidden_states, + temb, + image_only_indicator=image_only_indicator, + ) + + output_states = output_states + (hidden_states,) + + if exist_module_idx is not None and exist_module_idx == len(output_states) - 1: + return hidden_states, output_states + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class CrossAttnDownBlockSpatioTemporal(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + num_layers: int = 1, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + num_attention_heads: int = 1, + cross_attention_dim: int = 1280, + add_downsample: bool = True, + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * num_layers + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + SpatioTemporalResBlock( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=1e-6, + ) + ) + attentions.append( + TransformerSpatioTemporalModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block[i], + cross_attention_dim=cross_attention_dim, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=1, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + image_only_indicator: Optional[torch.Tensor] = None, + exist_module_idx: Optional[int] = None, + ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: + output_states = () + + blocks = list(zip(self.resnets, self.attentions)) + for resnet, attn in blocks: + if self.training and self.gradient_checkpointing: # TODO + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + **ckpt_kwargs, + ) + + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + return_dict=False, + )[0] + else: + hidden_states = resnet( + hidden_states, + temb, + image_only_indicator=image_only_indicator, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + return_dict=False, + )[0] + + output_states = output_states + (hidden_states,) + if exist_module_idx is not None and exist_module_idx == len(output_states) - 1: + return hidden_states, output_states + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class UpBlockSpatioTemporal(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + resolution_idx: Optional[int] = None, + num_layers: int = 1, + resnet_eps: float = 1e-6, + add_upsample: bool = True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + SpatioTemporalResBlock( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + image_only_indicator: Optional[torch.Tensor] = None, + enter_module_idx: Optional[int] = None, + ) -> torch.FloatTensor: + prv_f = [] + for idx, resnet in enumerate(self.resnets): + if enter_module_idx is not None and idx < enter_module_idx: + continue + + prv_f.append(hidden_states) + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + use_reentrant=False, + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + ) + else: + hidden_states = resnet( + hidden_states, + temb, + image_only_indicator=image_only_indicator, + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states, prv_f + + +class CrossAttnUpBlockSpatioTemporal(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + resolution_idx: Optional[int] = None, + num_layers: int = 1, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + resnet_eps: float = 1e-6, + num_attention_heads: int = 1, + cross_attention_dim: int = 1280, + add_upsample: bool = True, + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * num_layers + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + SpatioTemporalResBlock( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + ) + ) + attentions.append( + TransformerSpatioTemporalModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block[i], + cross_attention_dim=cross_attention_dim, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + image_only_indicator: Optional[torch.Tensor] = None, + enter_module_idx: Optional[int] = None, + ) -> torch.FloatTensor: + prv_f = [] + for idx, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)): + if enter_module_idx is not None and idx < enter_module_idx: + continue + + prv_f.append(hidden_states) + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: # TODO + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + image_only_indicator, + **ckpt_kwargs, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + return_dict=False, + )[0] + else: + hidden_states = resnet( + hidden_states, + temb, + image_only_indicator=image_only_indicator, + ) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + return_dict=False, + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states, prv_f diff --git a/ixformer_sdk/contrib/DeepCache/svd/unet_spatio_temporal_condition.py b/ixformer_sdk/contrib/DeepCache/svd/unet_spatio_temporal_condition.py new file mode 100644 index 00000000..abd6092e --- /dev/null +++ b/ixformer_sdk/contrib/DeepCache/svd/unet_spatio_temporal_condition.py @@ -0,0 +1,566 @@ +from dataclasses import dataclass +from typing import Dict, Optional, Tuple, Union + +import torch +import torch.nn as nn + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import UNet2DConditionLoadersMixin +from diffusers.utils import BaseOutput, logging +from diffusers.models.attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from diffusers.models.modeling_utils import ModelMixin + +from .unet_3d_blocks import UNetMidBlockSpatioTemporal, get_down_block, get_up_block + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class UNetSpatioTemporalConditionOutput(BaseOutput): + """ + The output of [`UNetSpatioTemporalConditionModel`]. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`): + The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. + """ + + sample: torch.FloatTensor = None + + +class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): + r""" + A conditional Spatio-Temporal UNet model that takes a noisy video frames, conditional state, and a timestep and returns a sample + shaped output. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Parameters: + sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + Height and width of input/output sample. + in_channels (`int`, *optional*, defaults to 8): Number of channels in the input sample. + out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. + down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`): + The tuple of downsample blocks to use. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`): + The tuple of upsample blocks to use. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + The tuple of output channels for each block. + addition_time_embed_dim: (`int`, defaults to 256): + Dimension to to encode the additional time ids. + projection_class_embeddings_input_dim (`int`, defaults to 768): + The dimension of the projection of encoded `added_time_ids`. + layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. + cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + The dimension of the cross attention features. + transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1): + The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for + [`~models.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal`], [`~models.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal`], + [`~models.unet_3d_blocks.UNetMidBlockSpatioTemporal`]. + num_attention_heads (`int`, `Tuple[int]`, defaults to `(5, 10, 10, 20)`): + The number of attention heads. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 8, + out_channels: int = 4, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlockSpatioTemporal", + "CrossAttnDownBlockSpatioTemporal", + "CrossAttnDownBlockSpatioTemporal", + "DownBlockSpatioTemporal", + ), + up_block_types: Tuple[str] = ( + "UpBlockSpatioTemporal", + "CrossAttnUpBlockSpatioTemporal", + "CrossAttnUpBlockSpatioTemporal", + "CrossAttnUpBlockSpatioTemporal", + ), + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + addition_time_embed_dim: int = 256, + projection_class_embeddings_input_dim: int = 768, + layers_per_block: Union[int, Tuple[int]] = 2, + cross_attention_dim: Union[int, Tuple[int]] = 1024, + transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, + num_attention_heads: Union[int, Tuple[int]] = (5, 10, 10, 20), + num_frames: int = 25, + ): + super().__init__() + + self.sample_size = sample_size + + # Check inputs + if len(down_block_types) != len(up_block_types): + raise ValueError( + f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." + ) + + if len(block_out_channels) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." + ) + + if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." + ) + + # input + self.conv_in = nn.Conv2d( + in_channels, + block_out_channels[0], + kernel_size=3, + padding=1, + ) + + # time + time_embed_dim = block_out_channels[0] * 4 + + self.time_proj = Timesteps(block_out_channels[0], True, downscale_freq_shift=0) + timestep_input_dim = block_out_channels[0] + + self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) + + self.add_time_proj = Timesteps(addition_time_embed_dim, True, downscale_freq_shift=0) + self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) + + self.down_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + if isinstance(num_attention_heads, int): + num_attention_heads = (num_attention_heads,) * len(down_block_types) + + if isinstance(cross_attention_dim, int): + cross_attention_dim = (cross_attention_dim,) * len(down_block_types) + + if isinstance(layers_per_block, int): + layers_per_block = [layers_per_block] * len(down_block_types) + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) + + blocks_time_embed_dim = time_embed_dim + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block[i], + transformer_layers_per_block=transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + temb_channels=blocks_time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=1e-5, + cross_attention_dim=cross_attention_dim[i], + num_attention_heads=num_attention_heads[i], + resnet_act_fn="silu", + ) + self.down_blocks.append(down_block) + + # mid + self.mid_block = UNetMidBlockSpatioTemporal( + block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + transformer_layers_per_block=transformer_layers_per_block[-1], + cross_attention_dim=cross_attention_dim[-1], + num_attention_heads=num_attention_heads[-1], + ) + + # count how many layers upsample the images + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_num_attention_heads = list(reversed(num_attention_heads)) + reversed_layers_per_block = list(reversed(layers_per_block)) + reversed_cross_attention_dim = list(reversed(cross_attention_dim)) + reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block)) + + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=reversed_layers_per_block[i] + 1, + transformer_layers_per_block=reversed_transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=blocks_time_embed_dim, + add_upsample=add_upsample, + resnet_eps=1e-5, + resolution_idx=i, + cross_attention_dim=reversed_cross_attention_dim[i], + num_attention_heads=reversed_num_attention_heads[i], + resnet_act_fn="silu", + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-5) + self.conv_act = nn.SiLU() + + self.conv_out = nn.Conv2d( + block_out_channels[0], + out_channels, + kernel_size=3, + padding=1, + ) + + @property + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors( + name: str, + module: torch.nn.Module, + processors: Dict[str, AttentionProcessor], + ): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f"{name}.processor")) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor) + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + # Copied from diffusers.models.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking + def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None: + """ + Sets the attention processor to use [feed forward + chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). + + Parameters: + chunk_size (`int`, *optional*): + The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually + over each tensor of dim=`dim`. + dim (`int`, *optional*, defaults to `0`): + The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch) + or dim=1 (sequence length). + """ + if dim not in [0, 1]: + raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}") + + # By default chunk size is 1 + chunk_size = chunk_size or 1 + + def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int): + if hasattr(module, "set_chunk_feed_forward"): + module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) + + for child in module.children(): + fn_recursive_feed_forward(child, chunk_size, dim) + + for module in self.children(): + fn_recursive_feed_forward(module, chunk_size, dim) + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + added_time_ids: torch.Tensor, + cache_features: Optional[torch.Tensor] = None, + cache_branch: Optional[int] = None, + return_dict: bool = True, + ) -> Union[UNetSpatioTemporalConditionOutput, Tuple]: + r""" + The [`UNetSpatioTemporalConditionModel`] forward method. + + Args: + sample (`torch.FloatTensor`): + The noisy input tensor with the following shape `(batch, num_frames, channel, height, width)`. + timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. + encoder_hidden_states (`torch.FloatTensor`): + The encoder hidden states with shape `(batch, sequence_length, cross_attention_dim)`. + added_time_ids: (`torch.FloatTensor`): + The additional time ids with shape `(batch, num_additional_ids)`. These are encoded with sinusoidal + embeddings and added to the time embeddings. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] instead of a plain + tuple. + Returns: + [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] or `tuple`: + If `return_dict` is True, an [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] is returned, otherwise + a `tuple` is returned where the first element is the sample tensor. + """ + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + batch_size, num_frames = sample.shape[:2] + timesteps = timesteps.expand(batch_size) + + t_emb = self.time_proj(timesteps) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb) + + time_embeds = self.add_time_proj(added_time_ids.flatten()) + time_embeds = time_embeds.reshape((batch_size, -1)) + time_embeds = time_embeds.to(emb.dtype) + aug_emb = self.add_embedding(time_embeds) + emb = emb + aug_emb + + # Flatten the batch and frames dimensions + # sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width] + sample = sample.flatten(0, 1) + # Repeat the embeddings num_video_frames times + # emb: [batch, channels] -> [batch * frames, channels] + emb = emb.repeat_interleave(num_frames, dim=0) + # encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels] + encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0) + + # 2. pre-process + sample = self.conv_in(sample) + + image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device) + + # Branch: 4 down_blocks, each with 3 skip connections. Here we ignore the first skip branch, whose computations only has up_blocks but without down_blocks. + if cache_branch is not None: + each_module_num = len(self.down_blocks[0].resnets) + 1 + down_cache_block_idx = cache_branch // each_module_num + down_cache_module_idx = cache_branch % each_module_num + + up_cache_block_idx = len(self.up_blocks) - 1 - down_cache_block_idx + up_cache_module_idx = 1 - down_cache_module_idx + if down_cache_module_idx == each_module_num - 1: + up_cache_block_idx -= 1 + up_cache_module_idx = 2 + + if cache_features is not None: + # 3. down + down_block_res_samples = (sample,) + for block_id, downsample_block in enumerate(self.down_blocks): + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + exist_module_idx=down_cache_module_idx if down_cache_block_idx == block_id else None + ) + else: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + image_only_indicator=image_only_indicator, + exist_module_idx=down_cache_module_idx if down_cache_block_idx == block_id else None + ) + + down_block_res_samples += res_samples + if down_cache_block_idx == block_id: + break + + # 4. no mid + sample = cache_features + + # 5. up + for i, upsample_block in enumerate(self.up_blocks): + if i < up_cache_block_idx: + continue + + if i == up_cache_block_idx: + trunc_res_samples_len = len(upsample_block.resnets) - up_cache_module_idx + else: + trunc_res_samples_len = len(upsample_block.resnets) + + res_samples = down_block_res_samples[-trunc_res_samples_len :] + down_block_res_samples = down_block_res_samples[: -trunc_res_samples_len] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + sample, _ = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + enter_module_idx=up_cache_module_idx if i == up_cache_block_idx else None + ) + else: + sample, _ = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + image_only_indicator=image_only_indicator, + enter_module_idx=up_cache_module_idx if i == up_cache_block_idx else None + ) + else: + # 3. down + down_block_res_samples = (sample,) + for downsample_block in self.down_blocks: + if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + ) + else: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + image_only_indicator=image_only_indicator, + ) + + down_block_res_samples += res_samples + + # 4. mid + sample = self.mid_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + ) + + # 5. up + for i, upsample_block in enumerate(self.up_blocks): + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] + + if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: + sample, current_record_f = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + ) + else: + sample, current_record_f = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + image_only_indicator=image_only_indicator, + ) + + if cache_branch is not None and i == up_cache_block_idx: + cache_features = current_record_f[up_cache_module_idx] + + # 6. post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + # 7. Reshape back to original shape + sample = sample.reshape(batch_size, num_frames, *sample.shape[1:]) + + if not return_dict: + return (sample, cache_features) + + return UNetSpatioTemporalConditionOutput(sample=sample) diff --git a/ixformer_sdk/contrib/__init__.py b/ixformer_sdk/contrib/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/comfy/__init__.py b/ixformer_sdk/contrib/comfy/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/comfy/unet_model_wrapper.py b/ixformer_sdk/contrib/comfy/unet_model_wrapper.py new file mode 100644 index 00000000..422b17ea --- /dev/null +++ b/ixformer_sdk/contrib/comfy/unet_model_wrapper.py @@ -0,0 +1,496 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +import ixformer.functions as ixf_F + + +def time_embed(t_emb, weight1, bias1, weight2, bias2): + # unet time_emd + # linear + silu + linear + emb = ixf_F.act_bias_mm( + t_emb, weight1, act_type="silu", bias=bias1, scale=1, trans_format="TN" + ) + emb = ixf_F.act_bias_mm( + emb, weight2, act_type="none", bias=bias2, scale=1, trans_format="TN" + ) + return emb + + +def ixf_layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05): + return ixf_F.layernorm(input, weight, bias, normalized_shape) + + +def ixf_pt_scaled_dot_product_attention( + query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False +): + if ( + not query.is_contiguous() + and query.transpose(1, 2).is_contiguous() + and key.transpose(1, 2).is_contiguous() + and value.transpose(1, 2).is_contiguous() + and attn_mask is None + ): + + batch_size, head_num, seq_len_q, head_dim = query.shape + _, _, seq_len_k, _ = key.shape + + query = query.transpose(1, 2).view(batch_size * seq_len_q, head_num, head_dim) + key = key.transpose(1, 2).view(batch_size * seq_len_k, head_num, head_dim) + value = value.transpose(1, 2).view(batch_size * seq_len_k, head_num, head_dim) + + cu_seqlens_q = torch.arange( + 0, + seq_len_q * (batch_size + 1), + seq_len_q, + dtype=torch.int32, + device=query.device, + ) + if seq_len_q == seq_len_k: + cu_seqlens_k = cu_seqlens_q + else: + cu_seqlens_k = torch.arange( + 0, + seq_len_k * (batch_size + 1), + seq_len_k, + dtype=torch.int32, + device=query.device, + ) + + res = ixf_F.flash_attn_varlen_func( + query, + key, + value, + cu_seqlens_q.int(), + cu_seqlens_k.int(), + seq_len_q, + seq_len_k, + ) + res = res.view(batch_size, seq_len_q, head_num, head_dim).transpose(1, 2) + return res + + if not query.is_contiguous(): + query = query.contiguous() + if not key.is_contiguous(): + key = key.contiguous() + if not value.is_contiguous(): + value = value.contiguous() + return ixf_F.scaled_dot_product_attention( + query, key, value, attn_mask=attn_mask, is_causal=is_causal + ) + + +class UnetIxformerFunction: + def __init__(self) -> None: + self.ixf_linear = ixf_F.linear + self.pt_linear = F.linear + self.pt_layer_norm = F.layer_norm + self.pt_scaled_dot_product_attention = F.scaled_dot_product_attention + + def __enter__(self): + F.linear = self.ixf_linear + F.layer_norm = ixf_layer_norm + F.scaled_dot_product_attention = ixf_pt_scaled_dot_product_attention + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + F.linear = self.pt_linear + F.layer_norm = self.pt_layer_norm + F.scaled_dot_product_attention = self.pt_scaled_dot_product_attention + if exc_tb is not None: + print(f"{exc_type} {exc_val}") + return False + return True + + +def ForwardWrapper(fun): + def wrap(*args, **kwargs): + with UnetIxformerFunction() as w: + return fun(*args, **kwargs) + + return wrap + + +class IxformerComfyWrapper(nn.Module): + def __init__(self): + super().__init__() + self.is_ixf_wrapper = True + + +class Conv2dNhwcWrapper(IxformerComfyWrapper): + def __init__(self, module): + super().__init__() + module.weight.data = module.weight.permute(0, 2, 3, 1).contiguous() + module.bias.data = module.bias.float() + self.weight = module.weight.data + self.bias = module.bias.data + self.stride = module.stride + self.padding = module.padding + self.dilation = module.dilation + self.groups = module.groups + + def forward(self, x): + h2 = ixf_F.conv2d( + x, + self.weight, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + ) + return h2 + + +class ResBlockNhwcWrapper(IxformerComfyWrapper): + def __init__(self, module) -> None: + super().__init__() + assert not module.updown + assert not module.use_scale_shift_norm + assert not module.skip_t_emb + assert not module.exchange_temb_dims + + if isinstance(module.skip_connection, nn.Identity): + self.skip_connection = module.skip_connection + elif get_class_name(module.skip_connection) == "Conv2d": + self.skip_connection = Conv2dNhwcWrapper(module.skip_connection) + else: + raise NotImplementedError( + f"ResBlockNhwcWrapper support Conv2d or nn.Identity, but got {module.skip_connection}" + ) + + self.in_layers = module.in_layers + self.out_layers = module.out_layers + self.emb_layers = module.emb_layers + self.in_layers_conv = Conv2dNhwcWrapper(module.in_layers[2]) + self.out_layers_conv = Conv2dNhwcWrapper(module.out_layers[3]) + + def forward(self, x, emb): + # x: nhwc + x1 = x + # print(x1.shape) + # fused group_norm silu + h = ixf_F.group_norm( + x1, # nchw->nhwc + self.in_layers[0].num_groups, + self.in_layers[0].weight, + self.in_layers[0].bias, + format=False, + act_type=1, + ) + h = self.in_layers_conv(h) + + emb_out = self.emb_layers(emb) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + + h = h + emb_out.permute(0, 2, 3, 1) + # print(h.shape) + h = ixf_F.group_norm( + h, + self.out_layers[0].num_groups, + self.out_layers[0].weight, + self.out_layers[0].bias, + format=False, + act_type=1, + ) + + h = self.out_layers[2](h) + h = self.out_layers_conv(h) + # TODO: support other skip_connection + return self.skip_connection(x) + h + + +class DownsampleNhwcWrapper(IxformerComfyWrapper): + def __init__(self, module) -> None: + # TODO: support avg_pool_nd + super().__init__() + assert module.use_conv + self.channels = module.channels + self.op = Conv2dNhwcWrapper(module.op) + + def forward(self, x): + assert x.shape[-1] == self.channels + return self.op(x) + + +def ffn_forward(self, x): + if get_class_name(self.net[0]) == "GEGLU": + net = self.net[1:] + geglu_net = self.net[0] + x = geglu_net.proj(x) + x = ixf_F.gelu_and_mul(x) + return net(x) + else: + return self.net(x) + + +# ComfyUI/comfy/ldm/modules/attention.py `class BasicTransformerBlock(nn.Module)` +def transformer_block_forward(self, x, context=None, transformer_options={}): + extra_options = {} + block = transformer_options.get("block", None) + block_index = transformer_options.get("block_index", 0) + transformer_patches = {} + transformer_patches_replace = {} + + for k in transformer_options: + if k == "patches": + transformer_patches = transformer_options[k] + elif k == "patches_replace": + transformer_patches_replace = transformer_options[k] + else: + extra_options[k] = transformer_options[k] + + extra_options["n_heads"] = self.n_heads + extra_options["dim_head"] = self.d_head + + if self.ff_in: + x_skip = x + x = self.ff_in(self.norm_in(x)) + if self.is_res: + x += x_skip + + n = self.norm1(x) + if self.disable_self_attn: + context_attn1 = context + else: + context_attn1 = None + value_attn1 = None + + if "attn1_patch" in transformer_patches: + patch = transformer_patches["attn1_patch"] + if context_attn1 is None: + context_attn1 = n + value_attn1 = context_attn1 + for p in patch: + n, context_attn1, value_attn1 = p( + n, context_attn1, value_attn1, extra_options + ) + + if block is not None: + transformer_block = (block[0], block[1], block_index) + else: + transformer_block = None + + attn1_replace_patch = transformer_patches_replace.get("attn1", {}) + block_attn1 = transformer_block + if block_attn1 not in attn1_replace_patch: + block_attn1 = block + + if block_attn1 in attn1_replace_patch: + if context_attn1 is None: + context_attn1 = n + value_attn1 = n + n = self.attn1.to_q(n) + context_attn1 = self.attn1.to_k(context_attn1) + value_attn1 = self.attn1.to_v(value_attn1) + n = attn1_replace_patch[block_attn1]( + n, context_attn1, value_attn1, extra_options + ) + n = self.attn1.to_out(n) + else: + n = self.attn1(n, context=context_attn1, value=value_attn1) + + if "attn1_output_patch" in transformer_patches: + patch = transformer_patches["attn1_output_patch"] + for p in patch: + n = p(n, extra_options) + + x += n + if "middle_patch" in transformer_patches: + patch = transformer_patches["middle_patch"] + for p in patch: + x = p(x, extra_options) + + if self.attn2 is not None: + n = self.norm2(x) + if self.switch_temporal_ca_to_sa: + context_attn2 = n + else: + context_attn2 = context + value_attn2 = None + if "attn2_patch" in transformer_patches: + patch = transformer_patches["attn2_patch"] + value_attn2 = context_attn2 + for p in patch: + n, context_attn2, value_attn2 = p( + n, context_attn2, value_attn2, extra_options + ) + + attn2_replace_patch = transformer_patches_replace.get("attn2", {}) + block_attn2 = transformer_block + if block_attn2 not in attn2_replace_patch: + block_attn2 = block + + if block_attn2 in attn2_replace_patch: + if value_attn2 is None: + value_attn2 = context_attn2 + n = self.attn2.to_q(n) + context_attn2 = self.attn2.to_k(context_attn2) + value_attn2 = self.attn2.to_v(value_attn2) + n = attn2_replace_patch[block_attn2]( + n, context_attn2, value_attn2, extra_options + ) + n = self.attn2.to_out(n) + else: + n = self.attn2(n, context=context_attn2, value=value_attn2) + + if "attn2_output_patch" in transformer_patches: + patch = transformer_patches["attn2_output_patch"] + for p in patch: + n = p(n, extra_options) + + # x += n + # if self.is_res: + # x_skip = x + # x = self.ff(self.norm3(x)) + + x, x_skip = ixf_F.residual_layer_norm( + n, + self.norm3.normalized_shape, + self.norm3.weight, + self.norm3.bias, + x, + eps=self.norm3.eps, + is_post_ln=False, + ) + x = ffn_forward(self.ff, x) + + # x = ffn_forward(self.ff, self.norm3(x)) + if self.is_res: + x += x_skip + + return x + + +class SpatialTransformerNhwcWrapper(IxformerComfyWrapper): + def __init__(self, module): + super().__init__() + self.use_linear = module.use_linear + self.transformer_blocks = module.transformer_blocks + self.norm = module.norm + if not self.use_linear: + self.proj_in = Conv2dNhwcWrapper(module.proj_in) + self.proj_out = Conv2dNhwcWrapper(module.proj_out) + else: + self.proj_in = module.proj_in + self.proj_out = module.proj_out + + @ForwardWrapper + def forward(self, x, context=None, transformer_options={}): + # note: if no context is given, cross-attention defaults to self-attention + if not isinstance(context, list): + context = [context] * len(self.transformer_blocks) + + b, h, w, c = x.shape + x_in = x + + # group_norm + x = ixf_F.group_norm( + x, + self.norm.num_groups, + self.norm.weight, + self.norm.bias, + format=False, + ) + # conv2d + if not self.use_linear: + x = self.proj_in(x) + # n,(hw),c + x = x.view(x.shape[0], -1, x.shape[-1]) + if self.use_linear: + x = self.proj_in(x) + + for i, block in enumerate(self.transformer_blocks): + transformer_options["block_index"] = i + # x = block(x, context=context[i], transformer_options=transformer_options) + x = transformer_block_forward( + block, x, context=context[i], transformer_options=transformer_options + ) + + if self.use_linear: + x = self.proj_out(x) + x = x.view(b, h, w, c) + if not self.use_linear: + x = self.proj_out(x) + return x + x_in + + +class UpsampleNhwcWrapper(IxformerComfyWrapper): + def __init__(self, module) -> None: + # TODO: support mhwc interpolate + super().__init__() + self.dims = module.dims + self.use_conv = module.use_conv + self.channels = module.channels + if self.use_conv: + self.conv = Conv2dNhwcWrapper(module.conv) + + def forward(self, x, output_shape=None): + # print("================== Upsample is running ==================") + assert x.shape[-1] == self.channels + assert len(x.shape) == 4 + + # nhwc -> nchw + if output_shape is not None: + assert len(output_shape) == 4 + output_shape = [ + output_shape[0], + output_shape[3], + output_shape[1], + output_shape[2], + ] + x = x.permute(0, 3, 1, 2).contiguous() + if self.dims == 3: + shape = [x.shape[2], x.shape[3] * 2, x.shape[4] * 2] + if output_shape is not None: + shape[1] = output_shape[3] + shape[2] = output_shape[4] + else: + shape = [x.shape[2] * 2, x.shape[3] * 2] + if output_shape is not None: + shape[0] = output_shape[2] + shape[1] = output_shape[3] + # TODO: interpolate 支持 nhwc， 去掉前后转置 + x = F.interpolate(x, size=shape, mode="nearest") + # nchw -> nhwc + x = x.permute(0, 2, 3, 1).contiguous() + if self.use_conv: + x = self.conv(x) + return x + + +unet_wrappers = { + "Conv2d": Conv2dNhwcWrapper, + "ResBlock": ResBlockNhwcWrapper, + "Downsample": DownsampleNhwcWrapper, + "SpatialTransformer": SpatialTransformerNhwcWrapper, + "Upsample": UpsampleNhwcWrapper, +} + + +def get_class_name(module): + return module.__class__.__name__ + + +def module_wrapper(module): + # 将原始的 module 封装为 nhwc 模式 + module_name = get_class_name(module) + assert ( + module_name == "TimestepEmbedSequential" + ), f"ixformer unet_model_wrapper only support 'TimestepEmbedSequential' now, but got {module_name}" + + num_sequential = len(module) + for idx_seq in range(num_sequential): + sub_module = module[idx_seq] + sub_module_name = get_class_name(sub_module) + # 判断模块是否已经封装 + if not getattr(sub_module, "is_ixf_wrapper", False): + if sub_module_name in unet_wrappers: + module[idx_seq].forward = unet_wrappers[sub_module_name]( + sub_module + ).forward + module[idx_seq].is_ixf_wrapper = True + else: + raise NotImplementedError(f"{sub_module_name} not support") + return module diff --git a/ixformer_sdk/contrib/flashinfer/__init__.py b/ixformer_sdk/contrib/flashinfer/__init__.py new file mode 100644 index 00000000..6dab014b --- /dev/null +++ b/ixformer_sdk/contrib/flashinfer/__init__.py @@ -0,0 +1,17 @@ +from .decode import BatchDecodeWithPagedKVCacheWrapper +from .prefill import ( + BatchPrefillWithPagedKVCacheWrapper, + BatchPrefillWithRaggedKVCacheWrapper, +) + + +def bmm_fp8(): + pass + + +def SegmentGEMMWrapper(): + pass + + +def bmm_fp8(): + pass diff --git a/ixformer_sdk/contrib/flashinfer/activation.py b/ixformer_sdk/contrib/flashinfer/activation.py new file mode 100644 index 00000000..d309ec9b --- /dev/null +++ b/ixformer_sdk/contrib/flashinfer/activation.py @@ -0,0 +1,29 @@ +import ixformer.inference.functions as ops +import torch + + +def gelu_and_mul(): + pass + + +def gelu_tanh_and_mul(): + pass + + +def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor: + r"""Fused SiLU and Mul operation. + + Parameters + ---------- + input: torch.Tensor + Input tensor, shape (..., 2 * hidden_size). + + out: Optional[torch.Tensor] + The the output tensor, if specified, the kernel will update this tensor inplace. + + Returns + ------- + output: torch.Tensor + Output tensor, shape (..., hidden_size). + """ + return ops.silu_and_mul(input=input, output=out) diff --git a/ixformer_sdk/contrib/flashinfer/cascade.py b/ixformer_sdk/contrib/flashinfer/cascade.py new file mode 100644 index 00000000..8f5e9c8a --- /dev/null +++ b/ixformer_sdk/contrib/flashinfer/cascade.py @@ -0,0 +1,2 @@ +def merge_state(): + pass diff --git a/ixformer_sdk/contrib/flashinfer/decode.py b/ixformer_sdk/contrib/flashinfer/decode.py new file mode 100644 index 00000000..e1b5f556 --- /dev/null +++ b/ixformer_sdk/contrib/flashinfer/decode.py @@ -0,0 +1,101 @@ +import math +from typing import Optional, Tuple, Union + +import ixformer.inference.functions as ops +import torch + + +def _grouped_size_compiled_for_decode_kernels( + num_qo_heads: int, num_kv_heads: int +) -> bool: + return (num_qo_heads // num_kv_heads) in [1, 2, 4, 8] + + +class BatchDecodeWithPagedKVCacheWrapper: + def __init__( + self, + float_workspace_buffer: torch.Tensor, + kv_layout: str = "NHD", + use_cuda_graph: bool = False, + use_tensor_cores: bool = False, + ) -> None: + pass + + def plan( + self, + indptr: torch.Tensor, + indices: torch.Tensor, + last_page_len: torch.Tensor, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + page_size: int, + # pos_encoding_mode: str = "NONE", + # window_left: int = -1, + # logits_soft_cap: Optional[float] = None, + data_type: Union[str, torch.dtype] = "float16", + q_data_type: Optional[Union[str, torch.dtype]] = None, + sm_scale: Optional[float] = None, + # rope_scale: Optional[float] = None, + # rope_theta: Optional[float] = None, + max_seqlen_q: int = None, + max_seqlen_k: int = None, + ) -> None: + self.indptr = indptr + self.indices = indices + self.last_page_len = last_page_len + self.num_qo_heads = num_qo_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + + assert page_size == 1 + + self.cu_seqlens_q = torch.ones_like(indptr) + self.cu_seqlens_q[0] = 0 + self.cu_seqlens_q = torch.cumsum(self.cu_seqlens_q, dim=0).int() + + self.cu_seqlens_k = indptr + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + + self.sm_scale = sm_scale + self.max_seqlen_q = max_seqlen_q + self.max_seqlen_k = max_seqlen_k + + begin_forward = plan + + def forward( + self, + q: torch.Tensor, + paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + pos_encoding_mode: str = "NONE", + q_scale: Optional[float] = None, + k_scale: Optional[float] = None, + v_scale: Optional[float] = None, + window_left: int = -1, + logits_soft_cap: Optional[float] = None, + sm_scale: Optional[float] = None, + rope_scale: Optional[float] = None, + rope_theta: Optional[float] = None, + ) -> torch.Tensor: + k_cache, v_cache = paged_kv_cache + + out = torch.empty_like(q) + + ops.paged_attention_flashinfer( + output=out, + query=q, + paged_kv_data=(k_cache.unsqueeze(1), v_cache.unsqueeze(1)), + paged_kv_indptr=self.indptr, + paged_kv_indices=self.indices, + paged_kv_last_page_len=self.last_page_len, + scale=self.sm_scale, + max_seq_len=self.max_seqlen_k, + kv_cache_format="NHD", + ) + + return out + + def end_forward(self) -> None: + r"""Warning: this function is deprecated and has no effect.""" + pass diff --git a/ixformer_sdk/contrib/flashinfer/norm.py b/ixformer_sdk/contrib/flashinfer/norm.py new file mode 100644 index 00000000..53ef36fd --- /dev/null +++ b/ixformer_sdk/contrib/flashinfer/norm.py @@ -0,0 +1,61 @@ +import ixformer.inference.functions as ops +import torch + + +def fused_add_rmsnorm( + input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6 +): + r"""Fused add root mean square normalization. + + Parameters + ---------- + input: torch.Tensor + Input tensor, shape (batch_size, hidden_size). + residual: torch.Tensor + Residual tensor, shape (batch_size, hidden_size). + weight: torch.Tensor + Weight tensor, shape (hidden_size,). + eps: float + Epsilon for numerical stability. + """ + return ops.residual_rms_norm( + input=input, + residual=residual, + weight=weight, + eps=eps, + ) + + +def gemma_fused_add_rmsnorm(): + pass + + +def gemma_rmsnorm(): + pass + + +def rmsnorm( + input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6 +) -> torch.Tensor: + r"""Root mean square normalization. + + Parameters + ---------- + input: torch.Tensor + Input tensor, shape (batch_size, hidden_size). + weight: torch.Tensor + Weight tensor, shape (hidden_size,). + eps: float + Epsilon for numerical stability. + + Returns + ------- + output: torch.Tensor + Normalized tensor, shape (batch_size, hidden_size). + """ + + return ops.rms_norm( + input=input, + weight=weight, + eps=eps, + ) diff --git a/ixformer_sdk/contrib/flashinfer/prefill.py b/ixformer_sdk/contrib/flashinfer/prefill.py new file mode 100644 index 00000000..a8a624a8 --- /dev/null +++ b/ixformer_sdk/contrib/flashinfer/prefill.py @@ -0,0 +1,113 @@ +import math +from typing import Optional, Tuple, Union + +import ixformer._C as ops +import torch + + +class BatchPrefillWithRaggedKVCacheWrapper: + def __init__( + self, + float_workspace_buffer: torch.Tensor, + kv_layout: str = "NHD", + ): + pass + + def plan( + self, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + max_seqlen_q: int, + max_seqlen_k: int, + # custom_mask: Optional[torch.Tensor] = None, + # packed_custom_mask: Optional[torch.Tensor] = None, + causal: bool = True, + # pos_encoding_mode: str = "NONE", + # allow_fp16_qk_reduction: bool = False, + # window_left: int = -1, + # logits_soft_cap: Optional[float] = None, + sm_scale: Optional[float] = None, + # rope_scale: Optional[float] = None, + # rope_theta: Optional[float] = None, + # q_data_type: str = "float16", + ) -> None: + batch_size = len(qo_indptr) - 1 + if len(kv_indptr) != batch_size + 1: + raise ValueError( + "The kv_indptr length should be equal to qk_indptr length." + ) + self._causal = causal + self._sm_scale = sm_scale + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + + self.cu_seqlens_q = qo_indptr + self.cu_seqlens_k = kv_indptr + self.num_qo_heads = num_qo_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + self.max_seqlen_q = max_seqlen_q + self.max_seqlen_k = max_seqlen_k + + begin_forward = plan + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool = True, + # pos_encoding_mode: str = "NONE", + # allow_fp16_qk_reduction: bool = False, + # window_left: int = -1, + logits_soft_cap: Optional[float] = None, + sm_scale: Optional[float] = None, + # rope_scale: Optional[float] = None, + # rope_theta: Optional[float] = None, + ) -> torch.Tensor: + r"""Warning: This function is deprecated, please use :meth:`run` instead.""" + + q = q.view(-1, self.num_qo_heads, self.head_dim) + k = k.view(-1, self.num_kv_heads, self.head_dim) + v = v.view(-1, self.num_kv_heads, self.head_dim) + + out = torch.empty_like(q) + + assert causal + assert ( + logits_soft_cap is None or logits_soft_cap == 0 + ), f"logits_soft_cap not supported, but got logits_soft_cap={logits_soft_cap}" + + ops.infer.ixinfer_flash_attn_unpad( + q, + k, + v, + out, + self.cu_seqlens_q, + self.cu_seqlens_k, + self.max_seqlen_q, + self.max_seqlen_k, + causal, + False, # need_lse =False + sm_scale, + False, + None, + ) + return out + + def end_forward(self) -> None: + r"""Warning: this function is deprecated and has no effect.""" + pass + + +class BatchPrefillWithPagedKVCacheWrapper: + def __init__( + self, + float_workspace_buffer: torch.Tensor, + kv_layout: str = "NHD", + use_cuda_graph: bool = False, + ) -> None: + pass diff --git a/ixformer_sdk/contrib/flashinfer/sampling.py b/ixformer_sdk/contrib/flashinfer/sampling.py new file mode 100644 index 00000000..47025152 --- /dev/null +++ b/ixformer_sdk/contrib/flashinfer/sampling.py @@ -0,0 +1,14 @@ +def min_p_sampling_from_probs(): + pass + + +def top_k_renorm_prob(): + pass + + +def top_k_top_p_sampling_from_probs(): + pass + + +def top_p_renorm_prob(): + pass diff --git a/ixformer_sdk/contrib/tgi/__init__.py b/ixformer_sdk/contrib/tgi/__init__.py new file mode 100644 index 00000000..c951b1c1 --- /dev/null +++ b/ixformer_sdk/contrib/tgi/__init__.py @@ -0,0 +1 @@ +from .fused_moe import fused_moe \ No newline at end of file diff --git a/ixformer_sdk/contrib/tgi/fused_moe.py b/ixformer_sdk/contrib/tgi/fused_moe.py new file mode 100644 index 00000000..f57b8c0a --- /dev/null +++ b/ixformer_sdk/contrib/tgi/fused_moe.py @@ -0,0 +1,430 @@ +import functools +import json +import os +from typing import Any, Dict, Optional, Tuple +from loguru import logger +import torch +import ixformer.inference.functions as ops + +CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "65536")) + +def fused_topk( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, +): + assert hidden_states.shape[0] == gating_output.shape[0], ( + "Number of tokens mismatch") + + M, _ = hidden_states.shape + + topk_weights = torch.empty(M, + topk, + dtype=torch.float32, + device=hidden_states.device) + topk_ids = torch.empty(M, + topk, + dtype=torch.int32, + device=hidden_states.device) + token_expert_indicies = torch.empty(M, + topk, + dtype=torch.int32, + device=hidden_states.device) + ops.vllm_moe_topk_softmax( + topk_weights, + topk_ids, + token_expert_indicies, + gating_output.float(), # TODO(woosuk): Optimize this. + ) + del token_expert_indicies # Not used. Will be used in the future. + + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + return topk_weights, topk_ids + +# This is used by the Deepseek-V2 model +def grouped_topk(hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: int = 0, + topk_group: int = 0): + + assert hidden_states.shape[0] == gating_output.shape[0], ( + "Number of tokens mismatch") + + scores = torch.softmax(gating_output, dim=-1) + num_token = scores.shape[0] + group_scores = scores.view(num_token, num_expert_group, + -1).max(dim=-1).values # [n, n_group] + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, + sorted=False)[1] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = group_mask.unsqueeze(-1).expand( + num_token, num_expert_group, + scores.shape[-1] // num_expert_group).reshape(num_token, -1) # [n, e] + tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e] + topk_weights, topk_ids = torch.topk(tmp_scores, + k=topk, + dim=-1, + sorted=False) + + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + return topk_weights, topk_ids + +def get_config_file_name(E: int, N: int, dtype: Optional[str]) -> str: + device_name = torch.cuda.get_device_name().replace(" ", "_") + dtype_selector = "" if not dtype else f",dtype={dtype}" + return f"E={E},N={N},device_name={device_name}{dtype_selector}.json" + +@functools.lru_cache +def get_moe_configs(E: int, N: int, + dtype: Optional[str]) -> Optional[Dict[int, Any]]: + """ + Return optimized configurations for the fused MoE kernel. + + The return value will be a dictionary that maps an irregular grid of + batch sizes to configurations of the fused_moe kernel. To evaluate the + kernel on a given batch size bs, the closest batch size in the grid should + be picked and the associated configuration chosen to invoke the kernel. + """ + + # First look up if an optimized configuration is available in the configs + # directory + json_file_name = get_config_file_name(E, N, dtype) + + config_file_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name) + if os.path.exists(config_file_path): + with open(config_file_path) as f: + logger.info("Using configuration from %s for MoE layer.", + config_file_path) + # If a configuration has been found, return it + return {int(key): val for key, val in json.load(f).items()} + + # If no optimized configuration is available, we will use the default + # configuration + return None + +def get_default_config( + M: int, + E: int, + N: int, + K: int, + topk: int, + dtype: Optional[str], +) -> Dict[str, int]: + config = { + 'BLOCK_SIZE_M': 64, + 'BLOCK_SIZE_N': 64, + 'BLOCK_SIZE_K': 32, + 'GROUP_SIZE_M': 8 + } + if M <= E: + config = { + 'BLOCK_SIZE_M': 16, + 'BLOCK_SIZE_N': 32, + 'BLOCK_SIZE_K': 64, + 'GROUP_SIZE_M': 1 + } + numel = M * topk + if numel <= 64: + config['BLOCK_SIZE_M'] = 32 + elif numel <= 1024: + config['BLOCK_SIZE_M'] = 64 + else: + config['BLOCK_SIZE_M'] = 256 + return config + +def try_get_optimal_moe_config( + w1_shape: Tuple[int, ...], + w2_shape: Tuple[int, ...], + top_k: int, + dtype: Optional[str], + M: int, + override_config: Optional[Dict[str, Any]] = None, +): + if override_config: + config = override_config + else: + # First try to load optimal config from the file + E, _, N = w2_shape + configs = get_moe_configs(E, N, dtype) + + if configs: + # If an optimal configuration map has been found, look up the + # optimal config + config = configs[min(configs.keys(), key=lambda x: abs(x - M))] + else: + # Else use the default config + config = get_default_config(M, E, N, w1_shape[2], top_k, dtype) + return config + +def moe_align_block_size( + topk_ids: torch.Tensor, block_size: int, + num_experts: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Aligns the token distribution across experts to be compatible with block + size for matrix multiplication. + + Parameters: + - topk_ids: A tensor of shape [total_tokens, top_k] representing the + top-k expert indices for each token. + - block_size: The block size used in block matrix multiplication. + - num_experts: The total number of experts. + + Returns: + - sorted_token_ids: A tensor containing the sorted token indices according + to their allocated expert. + - expert_ids: A tensor indicating the assigned expert index for each block. + - num_tokens_post_padded: The total number of tokens after padding, + ensuring divisibility by block_size. + + This function pads the number of tokens that each expert needs to process + so that it is divisible by block_size. + Padding ensures that during block matrix multiplication, the dimensions + align correctly. + + Example: + Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]], + block_size = 4, and num_experts = 4: + - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts, + with each expert needing to process 3 tokens. + - As block_size is 4, we pad 1 token for each expert. + - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3]. + - Then append padding tokens [12, 12, 12, 12] for each block. + - After sorting by expert index, we obtain token_ids + [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12]. + Tokens 12 are non-existent (padding) and are ignored in + the subsequent matrix multiplication. + - The padding ensures that the total number of tokens is now divisible + by block_size for proper block matrix operations. + """ + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) + sorted_ids = torch.empty((max_num_tokens_padded, ), + dtype=torch.int32, + device=topk_ids.device) + sorted_ids.fill_(topk_ids.numel()) + # max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size) + max_num_m_blocks = topk_ids.numel() + num_experts + expert_ids = torch.empty((max_num_m_blocks, ), + dtype=torch.int32, + device=topk_ids.device) + num_tokens_post_pad = torch.empty((1), + dtype=torch.int32, + device=topk_ids.device) + ops.vllm_moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids, + expert_ids, num_tokens_post_pad) + return sorted_ids, expert_ids, num_tokens_post_pad + +def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, + A_scale: Optional[torch.Tensor], + B_scale: Optional[torch.Tensor], + topk_weights: torch.Tensor, topk_ids: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + mul_routed_weight: bool, top_k: int, + config: Dict[str, Any], compute_type: torch.dtype, + use_fp8: bool) -> None: + ops.vllm_invoke_fused_moe_kernel(A, B, C, topk_weights, topk_ids, + sorted_token_ids,expert_ids, num_tokens_post_padded, + mul_routed_weight, top_k, config['BLOCK_SIZE_M']) + +def fused_experts(hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + inplace: bool = False, + override_config: Optional[Dict[str, Any]] = None, + use_fp8: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None): + # Check constraints. + assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" + assert w1.is_contiguous(), "Expert weights1 must be contiguous" + assert w2.is_contiguous(), "Expert weights2 must be contiguous" + assert hidden_states.dtype in [ + torch.float32, torch.float16, torch.bfloat16 + ] + + num_tokens, _ = hidden_states.shape + E, N, _ = w1.shape + # We execute the fused_moe kernel in chunks to circumvent this issue: + # https://github.com/vllm-project/vllm/issues/5938 + M = min(num_tokens, CHUNK_SIZE) + + get_config_func = functools.partial( + try_get_optimal_moe_config, + w1.shape, + w2.shape, + topk_ids.shape[1], + "float8" if use_fp8 else None, + override_config=override_config, + ) + + config = get_config_func(M) + + intermediate_cache1 = torch.empty((M, topk_ids.shape[1], N), + device=hidden_states.device, + dtype=hidden_states.dtype) + intermediate_cache2 = torch.empty((M * topk_ids.shape[1], N // 2), + device=hidden_states.device, + dtype=hidden_states.dtype) + intermediate_cache3 = torch.empty((M, topk_ids.shape[1], w2.shape[1]), + device=hidden_states.device, + dtype=hidden_states.dtype) + + compute_type = (torch.bfloat16 + if hidden_states.dtype == torch.bfloat16 else torch.float16) + + if inplace: + out_hidden_states = hidden_states + else: + out_hidden_states = torch.empty_like(hidden_states) + + for chunk in range((num_tokens // CHUNK_SIZE) + 1): + begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE, + min((chunk + 1) * CHUNK_SIZE, + num_tokens)) + curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx] + tokens_in_chunk, _ = curr_hidden_states.shape + + if tokens_in_chunk == 0: + break + + if tokens_in_chunk < CHUNK_SIZE and chunk > 0: + # Adjust the intermediate cache size and config for the last + # chunk. Note that in most cases we only have one chunk + # so the cache size and config are already set correctly and + # do not need to be adjusted. + intermediate_cache1 = intermediate_cache1[:tokens_in_chunk] + intermediate_cache2 = intermediate_cache2[:tokens_in_chunk] + intermediate_cache3 = intermediate_cache3[:tokens_in_chunk] + config = get_config_func(tokens_in_chunk) + + curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx] + curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx] + + sorted_token_ids, expert_ids, num_tokens_post_padded = ( + moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'], E)) + + invoke_fused_moe_kernel(curr_hidden_states, + w1, + intermediate_cache1, + a1_scale, + w1_scale, + curr_topk_weights, + curr_topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + False, + topk_ids.shape[1], + config, + compute_type=compute_type, + use_fp8=use_fp8) + + ops.silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2) + + invoke_fused_moe_kernel(intermediate_cache2, + w2, + intermediate_cache3, + a2_scale, + w2_scale, + curr_topk_weights, + curr_topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + True, + 1, + config, + compute_type=compute_type, + use_fp8=use_fp8) + + torch.sum(intermediate_cache3.view(*intermediate_cache3.shape), + dim=1, + out=out_hidden_states[begin_chunk_idx:end_chunk_idx]) + return out_hidden_states + +def fused_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + inplace: bool = False, + override_config: Optional[Dict[str, Any]] = None, + use_grouped_topk: bool = False, + num_expert_group: Optional[int] = None, + topk_group: Optional[int] = None, + use_fp8: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + a1_scale: Optional[torch.Tensor] = None, + a2_scale: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """ + This function computes a Mixture of Experts (MoE) layer using two sets of + weights, w1 and w2, and top-k gating mechanism. + + Parameters: + - hidden_states (torch.Tensor): The input tensor to the MoE layer. + - w1 (torch.Tensor): The first set of expert weights. + - w2 (torch.Tensor): The second set of expert weights. + - gating_output (torch.Tensor): The output of the gating operation + (before softmax). + - topk (int): The number of top-k experts to select. + - renormalize (bool): If True, renormalize the top-k weights to sum to 1. + - inplace (bool): If True, perform the operation in-place. + Defaults to False. + - override_config (Optional[Dict[str, Any]]): Optional override + for the kernel configuration. + - num_expert_group: Optional[int]: additional parameter for grouped_topk + - topk_group: Optional[int]: additional parameter for grouped_topk + - use_grouped_topk: If True, use grouped_topk instead of fused_topk + note: Deepseekv2 model uses grouped_topk + - use_fp8 (bool): If True, use fp8 arithmetic to compute the inner + products for w1 and w2. Defaults to False. + - w1_scale (Optional[torch.Tensor]): Optional scale to be used for + w1. + - w2_scale (Optional[torch.Tensor]): Optional scale to be used for + w2. + + Returns: + - torch.Tensor: The output tensor after applying the MoE layer. + """ + # Check constraints. + assert gating_output.shape[1] == w1.shape[0], "Number of experts mismatch" + + if use_grouped_topk: + assert num_expert_group is not None and topk_group is not None + topk_weights, topk_ids = grouped_topk(hidden_states, gating_output, + topk, renormalize, + num_expert_group, topk_group) + else: + topk_weights, topk_ids = fused_topk(hidden_states, gating_output, topk, + renormalize) + + return fused_experts(hidden_states, + w1, + w2, + topk_weights, + topk_ids, + inplace=inplace, + override_config=override_config, + use_fp8=use_fp8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale) \ No newline at end of file diff --git a/ixformer_sdk/contrib/transformers/__init__.py b/ixformer_sdk/contrib/transformers/__init__.py new file mode 100644 index 00000000..c70d26db --- /dev/null +++ b/ixformer_sdk/contrib/transformers/__init__.py @@ -0,0 +1,2 @@ +from .models.bert.modeling_bert import BertForQuestionAnswering +from .models.t5.modeling_t5 import T5ForConditionalGeneration diff --git a/ixformer_sdk/contrib/transformers/models/__init__.py b/ixformer_sdk/contrib/transformers/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/transformers/models/bert/__init__.py b/ixformer_sdk/contrib/transformers/models/bert/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/transformers/models/bert/configuration_bert.py b/ixformer_sdk/contrib/transformers/models/bert/configuration_bert.py new file mode 100644 index 00000000..1db36391 --- /dev/null +++ b/ixformer_sdk/contrib/transformers/models/bert/configuration_bert.py @@ -0,0 +1,150 @@ +# coding=utf-8 +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# 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. +"""BERT model configuration""" + +from collections import OrderedDict +from typing import Mapping + +from transformers.configuration_utils import PretrainedConfig +from transformers.onnx import OnnxConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +class BertConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`BertModel`] or a [`TFBertModel`]. It is used to + instantiate a BERT model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the BERT + [google-bert/bert-base-uncased](https://huggingface.co/google-bert/bert-base-uncased) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 30522): + Vocabulary size of the BERT model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`BertModel`] or [`TFBertModel`]. + hidden_size (`int`, *optional*, defaults to 768): + Dimensionality of the encoder layers and the pooler layer. + num_hidden_layers (`int`, *optional*, defaults to 12): + Number of hidden layers in the Transformer encoder. + num_attention_heads (`int`, *optional*, defaults to 12): + Number of attention heads for each attention layer in the Transformer encoder. + intermediate_size (`int`, *optional*, defaults to 3072): + Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder. + hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`): + The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, + `"relu"`, `"silu"` and `"gelu_new"` are supported. + hidden_dropout_prob (`float`, *optional*, defaults to 0.1): + The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. + attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1): + The dropout ratio for the attention probabilities. + max_position_embeddings (`int`, *optional*, defaults to 512): + The maximum sequence length that this model might ever be used with. Typically set this to something large + just in case (e.g., 512 or 1024 or 2048). + type_vocab_size (`int`, *optional*, defaults to 2): + The vocabulary size of the `token_type_ids` passed when calling [`BertModel`] or [`TFBertModel`]. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + layer_norm_eps (`float`, *optional*, defaults to 1e-12): + The epsilon used by the layer normalization layers. + position_embedding_type (`str`, *optional*, defaults to `"absolute"`): + Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For + positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to + [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155). + For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models + with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658). + is_decoder (`bool`, *optional*, defaults to `False`): + Whether the model is used as a decoder or not. If `False`, the model is used as an encoder. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + classifier_dropout (`float`, *optional*): + The dropout ratio for the classification head. + + Examples: + + ```python + >>> from transformers import BertConfig, BertModel + + >>> # Initializing a BERT google-bert/bert-base-uncased style configuration + >>> configuration = BertConfig() + + >>> # Initializing a model (with random weights) from the google-bert/bert-base-uncased style configuration + >>> model = BertModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "bert" + + def __init__( + self, + vocab_size=30522, + hidden_size=768, + num_hidden_layers=12, + num_attention_heads=12, + intermediate_size=3072, + hidden_act="gelu", + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + max_position_embeddings=512, + type_vocab_size=2, + initializer_range=0.02, + layer_norm_eps=1e-12, + pad_token_id=0, + position_embedding_type="absolute", + use_cache=True, + classifier_dropout=None, + **kwargs, + ): + super().__init__(pad_token_id=pad_token_id, **kwargs) + + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.hidden_dropout_prob = hidden_dropout_prob + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.max_position_embeddings = max_position_embeddings + self.type_vocab_size = type_vocab_size + self.initializer_range = initializer_range + self.layer_norm_eps = layer_norm_eps + self.position_embedding_type = position_embedding_type + self.use_cache = use_cache + self.classifier_dropout = classifier_dropout + + +class BertOnnxConfig(OnnxConfig): + @property + def inputs(self) -> Mapping[str, Mapping[int, str]]: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"} + else: + dynamic_axis = {0: "batch", 1: "sequence"} + return OrderedDict( + [ + ("input_ids", dynamic_axis), + ("attention_mask", dynamic_axis), + ("token_type_ids", dynamic_axis), + ] + ) diff --git a/ixformer_sdk/contrib/transformers/models/bert/modeling_bert.py b/ixformer_sdk/contrib/transformers/models/bert/modeling_bert.py new file mode 100644 index 00000000..2187a630 --- /dev/null +++ b/ixformer_sdk/contrib/transformers/models/bert/modeling_bert.py @@ -0,0 +1,2145 @@ +# coding=utf-8 +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# 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. +"""PyTorch BERT model.""" + +import math +import os +import warnings +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import ixformer.inference.functions as ops +import torch +import torch.utils.checkpoint +from packaging import version +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from transformers.activations import ACT2FN +from transformers.generation import GenerationMixin +from transformers.modeling_attn_mask_utils import ( + _prepare_4d_attention_mask_for_sdpa, + _prepare_4d_causal_attention_mask_for_sdpa, +) +from transformers.modeling_outputs import ( + BaseModelOutputWithPastAndCrossAttentions, + BaseModelOutputWithPoolingAndCrossAttentions, + CausalLMOutputWithCrossAttentions, + MaskedLMOutput, + MultipleChoiceModelOutput, + NextSentencePredictorOutput, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.pytorch_utils import ( + apply_chunking_to_forward, + find_pruneable_heads_and_indices, + prune_linear_layer, +) +from transformers.utils import ( + ModelOutput, + add_code_sample_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + get_torch_version, + logging, + replace_return_docstrings, +) + +from .configuration_bert import BertConfig + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "google-bert/bert-base-uncased" +_CONFIG_FOR_DOC = "BertConfig" + +# TokenClassification docstring +_CHECKPOINT_FOR_TOKEN_CLASSIFICATION = ( + "dbmdz/bert-large-cased-finetuned-conll03-english" +) +_TOKEN_CLASS_EXPECTED_OUTPUT = "['O', 'I-ORG', 'I-ORG', 'I-ORG', 'O', 'O', 'O', 'O', 'O', 'I-LOC', 'O', 'I-LOC', 'I-LOC'] " +_TOKEN_CLASS_EXPECTED_LOSS = 0.01 + +# QuestionAnswering docstring +_CHECKPOINT_FOR_QA = "deepset/bert-base-cased-squad2" +_QA_EXPECTED_OUTPUT = "'a nice puppet'" +_QA_EXPECTED_LOSS = 7.41 +_QA_TARGET_START_INDEX = 14 +_QA_TARGET_END_INDEX = 15 + +# SequenceClassification docstring +_CHECKPOINT_FOR_SEQUENCE_CLASSIFICATION = "textattack/bert-base-uncased-yelp-polarity" +_SEQ_CLASS_EXPECTED_OUTPUT = "'LABEL_1'" +_SEQ_CLASS_EXPECTED_LOSS = 0.01 + + +def load_tf_weights_in_bert(model, config, tf_checkpoint_path): + """Load tf checkpoints in a pytorch model.""" + try: + import re + + import numpy as np + import tensorflow as tf + except ImportError: + logger.error( + "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " + "https://www.tensorflow.org/install/ for installation instructions." + ) + raise + tf_path = os.path.abspath(tf_checkpoint_path) + logger.info(f"Converting TensorFlow checkpoint from {tf_path}") + # Load weights from TF model + init_vars = tf.train.list_variables(tf_path) + names = [] + arrays = [] + for name, shape in init_vars: + logger.info(f"Loading TF weight {name} with shape {shape}") + array = tf.train.load_variable(tf_path, name) + names.append(name) + arrays.append(array) + + for name, array in zip(names, arrays): + name = name.split("/") + # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v + # which are not required for using pretrained model + if any( + n + in [ + "adam_v", + "adam_m", + "AdamWeightDecayOptimizer", + "AdamWeightDecayOptimizer_1", + "global_step", + ] + for n in name + ): + logger.info(f"Skipping {'/'.join(name)}") + continue + pointer = model + for m_name in name: + if re.fullmatch(r"[A-Za-z]+_\d+", m_name): + scope_names = re.split(r"_(\d+)", m_name) + else: + scope_names = [m_name] + if scope_names[0] == "kernel" or scope_names[0] == "gamma": + pointer = getattr(pointer, "weight") + elif scope_names[0] == "output_bias" or scope_names[0] == "beta": + pointer = getattr(pointer, "bias") + elif scope_names[0] == "output_weights": + pointer = getattr(pointer, "weight") + elif scope_names[0] == "squad": + pointer = getattr(pointer, "classifier") + else: + try: + pointer = getattr(pointer, scope_names[0]) + except AttributeError: + logger.info(f"Skipping {'/'.join(name)}") + continue + if len(scope_names) >= 2: + num = int(scope_names[1]) + pointer = pointer[num] + if m_name[-11:] == "_embeddings": + pointer = getattr(pointer, "weight") + elif m_name == "kernel": + array = np.transpose(array) + try: + if pointer.shape != array.shape: + raise ValueError( + f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched" + ) + except ValueError as e: + e.args += (pointer.shape, array.shape) + raise + logger.info(f"Initialize PyTorch weight {name}") + pointer.data = torch.from_numpy(array) + return model + + +class BertEmbeddings(nn.Module): + """Construct the embeddings from word, position and token_type embeddings.""" + + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding( + config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id + ) + self.position_embeddings = nn.Embedding( + config.max_position_embeddings, config.hidden_size + ) + self.token_type_embeddings = nn.Embedding( + config.type_vocab_size, config.hidden_size + ) + + # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load + # any TensorFlow checkpoint file + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.position_embedding_type = getattr( + config, "position_embedding_type", "absolute" + ) + self.register_buffer( + "position_ids", + torch.arange(config.max_position_embeddings).expand((1, -1)), + persistent=False, + ) + self.register_buffer( + "token_type_ids", + torch.zeros(self.position_ids.size(), dtype=torch.long), + persistent=False, + ) + self.eps = config.layer_norm_eps + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + token_type_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + past_key_values_length: int = 0, + ) -> torch.Tensor: + assert input_ids is not None + assert token_type_ids is not None + assert position_ids is not None + + return ops.bert_embedding( + self.word_embeddings.weight, + self.position_embeddings.weight, + self.token_type_embeddings.weight, + self.LayerNorm.weight, + self.LayerNorm.bias, + input_ids, + position_ids, + token_type_ids, + self.eps, + ) + + +class BertSelfAttention(nn.Module): + def __init__(self, config, position_embedding_type=None): + super().__init__() + if config.hidden_size % config.num_attention_heads != 0 and not hasattr( + config, "embedding_size" + ): + raise ValueError( + f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " + f"heads ({config.num_attention_heads})" + ) + + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + self.position_embedding_type = position_embedding_type or getattr( + config, "position_embedding_type", "absolute" + ) + if ( + self.position_embedding_type == "relative_key" + or self.position_embedding_type == "relative_key_query" + ): + self.max_position_embeddings = config.max_position_embeddings + self.distance_embedding = nn.Embedding( + 2 * config.max_position_embeddings - 1, self.attention_head_size + ) + + self.is_decoder = config.is_decoder + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + ( + self.num_attention_heads, + self.attention_head_size, + ) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.Tensor]: + mixed_query_layer = self.query(hidden_states) + + # If this is instantiated as a cross-attention module, the keys + # and values come from an encoder; the attention mask needs to be + # such that the encoder's padding tokens are not attended to. + is_cross_attention = encoder_hidden_states is not None + + if is_cross_attention and past_key_value is not None: + # reuse k,v, cross_attentions + key_layer = past_key_value[0] + value_layer = past_key_value[1] + attention_mask = encoder_attention_mask + elif is_cross_attention: + key_layer = self.transpose_for_scores(self.key(encoder_hidden_states)) + value_layer = self.transpose_for_scores(self.value(encoder_hidden_states)) + attention_mask = encoder_attention_mask + elif past_key_value is not None: + key_layer = self.transpose_for_scores(self.key(hidden_states)) + value_layer = self.transpose_for_scores(self.value(hidden_states)) + key_layer = torch.cat([past_key_value[0], key_layer], dim=2) + value_layer = torch.cat([past_key_value[1], value_layer], dim=2) + else: + key_layer = self.transpose_for_scores(self.key(hidden_states)) + value_layer = self.transpose_for_scores(self.value(hidden_states)) + + query_layer = self.transpose_for_scores(mixed_query_layer) + + use_cache = past_key_value is not None + if self.is_decoder: + # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # Further calls to cross_attention layer can then reuse all cross-attention + # key/value_states (first "if" case) + # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # all previous decoder key/value_states. Further calls to uni-directional self-attention + # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) + # if encoder bi-directional self-attention `past_key_value` is always `None` + past_key_value = (key_layer, value_layer) + + # Take the dot product between "query" and "key" to get the raw attention scores. + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + + if ( + self.position_embedding_type == "relative_key" + or self.position_embedding_type == "relative_key_query" + ): + query_length, key_length = query_layer.shape[2], key_layer.shape[2] + if use_cache: + position_ids_l = torch.tensor( + key_length - 1, dtype=torch.long, device=hidden_states.device + ).view(-1, 1) + else: + position_ids_l = torch.arange( + query_length, dtype=torch.long, device=hidden_states.device + ).view(-1, 1) + position_ids_r = torch.arange( + key_length, dtype=torch.long, device=hidden_states.device + ).view(1, -1) + distance = position_ids_l - position_ids_r + + positional_embedding = self.distance_embedding( + distance + self.max_position_embeddings - 1 + ) + positional_embedding = positional_embedding.to( + dtype=query_layer.dtype + ) # fp16 compatibility + + if self.position_embedding_type == "relative_key": + relative_position_scores = torch.einsum( + "bhld,lrd->bhlr", query_layer, positional_embedding + ) + attention_scores = attention_scores + relative_position_scores + elif self.position_embedding_type == "relative_key_query": + relative_position_scores_query = torch.einsum( + "bhld,lrd->bhlr", query_layer, positional_embedding + ) + relative_position_scores_key = torch.einsum( + "bhrd,lrd->bhlr", key_layer, positional_embedding + ) + attention_scores = ( + attention_scores + + relative_position_scores_query + + relative_position_scores_key + ) + + attention_scores = attention_scores / math.sqrt(self.attention_head_size) + if attention_mask is not None: + # Apply the attention mask is (precomputed for all layers in BertModel forward() function) + attention_scores = attention_scores + attention_mask + + # Normalize the attention scores to probabilities. + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + # Mask heads if we want to + if head_mask is not None: + attention_probs = attention_probs * head_mask + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + + outputs = ( + (context_layer, attention_probs) if output_attentions else (context_layer,) + ) + + if self.is_decoder: + outputs = outputs + (past_key_value,) + return outputs + + +class BertSdpaSelfAttention(BertSelfAttention): + def __init__(self, config, position_embedding_type=None): + super().__init__(config, position_embedding_type=position_embedding_type) + self.dropout_prob = config.attention_probs_dropout_prob + self.require_contiguous_qkv = version.parse( + get_torch_version() + ) < version.parse("2.2.0") + + # Adapted from BertSelfAttention + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.Tensor]: + if ( + self.position_embedding_type != "absolute" + or output_attentions + or head_mask is not None + ): + # TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once implemented. + logger.warning_once( + "BertSdpaSelfAttention is used but `torch.nn.functional.scaled_dot_product_attention` does not support " + "non-absolute `position_embedding_type` or `output_attentions=True` or `head_mask`. Falling back to " + "the manual attention implementation, but specifying the manual implementation will be required from " + "Transformers version v5.0.0 onwards. This warning can be removed using the argument " + '`attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states, + attention_mask, + head_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + output_attentions, + ) + + bsz, tgt_len, _ = hidden_states.size() + + query_layer = self.transpose_for_scores(self.query(hidden_states)) + + # If this is instantiated as a cross-attention module, the keys and values come from an encoder; the attention + # mask needs to be such that the encoder's padding tokens are not attended to. + is_cross_attention = encoder_hidden_states is not None + + current_states = encoder_hidden_states if is_cross_attention else hidden_states + attention_mask = ( + encoder_attention_mask if is_cross_attention else attention_mask + ) + + # Check `seq_length` of `past_key_value` == `len(current_states)` to support prefix tuning + if ( + is_cross_attention + and past_key_value + and past_key_value[0].shape[2] == current_states.shape[1] + ): + key_layer, value_layer = past_key_value + else: + key_layer = self.transpose_for_scores(self.key(current_states)) + value_layer = self.transpose_for_scores(self.value(current_states)) + if past_key_value is not None and not is_cross_attention: + key_layer = torch.cat([past_key_value[0], key_layer], dim=2) + value_layer = torch.cat([past_key_value[1], value_layer], dim=2) + + if self.is_decoder: + # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # Further calls to cross_attention layer can then reuse all cross-attention + # key/value_states (first "if" case) + # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # all previous decoder key/value_states. Further calls to uni-directional self-attention + # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) + # if encoder bi-directional self-attention `past_key_value` is always `None` + past_key_value = (key_layer, value_layer) + + # SDPA with memory-efficient backend is broken in torch==2.1.2 when using non-contiguous inputs and a custom + # attn_mask, so we need to call `.contiguous()` here. This was fixed in torch==2.2.0. + # Reference: https://github.com/pytorch/pytorch/issues/112577 + if ( + self.require_contiguous_qkv + and query_layer.device.type == "cuda" + and attention_mask is not None + ): + query_layer = query_layer.contiguous() + key_layer = key_layer.contiguous() + value_layer = value_layer.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create + # a causal mask in case tgt_len == 1. + is_causal = ( + True + if self.is_decoder + and not is_cross_attention + and attention_mask is None + and tgt_len > 1 + else False + ) + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_layer, + key_layer, + value_layer, + attn_mask=attention_mask, + dropout_p=self.dropout_prob if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, tgt_len, self.all_head_size) + + outputs = (attn_output,) + if self.is_decoder: + outputs = outputs + (past_key_value,) + return outputs + + +class BertSelfOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.eps = config.layer_norm_eps + + def forward( + self, hidden_states: torch.Tensor, input_tensor: torch.Tensor + ) -> torch.Tensor: + hidden_states = ops.linear(hidden_states, self.dense.weight, self.dense.bias) + hidden_states = ops.bert_add_norm( + input=hidden_states, + residual=input_tensor, + ln_weight=self.LayerNorm.weight, + ln_bias=self.LayerNorm.bias, + epsilon=self.eps, + ) + return hidden_states + + +BERT_SELF_ATTENTION_CLASSES = { + "eager": BertSelfAttention, + "sdpa": BertSdpaSelfAttention, +} + + +class BertAttention(nn.Module): + def __init__(self, config, position_embedding_type=None): + super().__init__() + self.self = BERT_SELF_ATTENTION_CLASSES[config._attn_implementation]( + config, position_embedding_type=position_embedding_type + ) + self.output = BertSelfOutput(config) + self.pruned_heads = set() + self.qkv_weight = None + self.qkv_bias = None + + def prune_heads(self, heads): + if len(heads) == 0: + return + heads, index = find_pruneable_heads_and_indices( + heads, + self.self.num_attention_heads, + self.self.attention_head_size, + self.pruned_heads, + ) + + # Prune linear layers + self.self.query = prune_linear_layer(self.self.query, index) + self.self.key = prune_linear_layer(self.self.key, index) + self.self.value = prune_linear_layer(self.self.value, index) + self.output.dense = prune_linear_layer(self.output.dense, index, dim=1) + + # Update hyper params and store pruned heads + self.self.num_attention_heads = self.self.num_attention_heads - len(heads) + self.self.all_head_size = ( + self.self.attention_head_size * self.self.num_attention_heads + ) + self.pruned_heads = self.pruned_heads.union(heads) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + output_attentions: Optional[bool] = False, + cu_seq_lens: Optional[torch.Tensor] = None, + enable_unpad: Optional[bool] = False, + max_seq_len: Optional[int] = None, + ) -> Tuple[torch.Tensor]: + if self.qkv_weight is None: + self.qkv_weight = torch.cat( + [self.self.query.weight, self.self.key.weight, self.self.value.weight], + dim=0, + ) + + self.qkv_bias = torch.cat( + [self.self.query.bias, self.self.key.bias, self.self.value.bias], dim=0 + ) + del self.self.query + del self.self.key + del self.self.value + qkv = ops.linear(hidden_states, self.qkv_weight, self.qkv_bias) + q, k, v = torch.chunk(qkv, 3, dim=-1) + + num_heads = self.self.num_attention_heads + head_size = self.self.attention_head_size + + if not enable_unpad: + + bs, seq_len, hidden_size = q.shape + + q = ( + q.view(bs, seq_len, num_heads, head_size) + .permute(0, 2, 1, 3) + .contiguous() + ) + k = ( + k.view(bs, seq_len, num_heads, head_size) + .permute(0, 2, 1, 3) + .contiguous() + ) + v = ( + v.view(bs, seq_len, num_heads, head_size) + .permute(0, 2, 1, 3) + .contiguous() + ) + + self_outputs = ops.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + ) + self_outputs = (self_outputs,) + attention_output = ( + self_outputs[0] + .permute(0, 2, 1, 3) + .contiguous() + .view(bs, seq_len, hidden_size) + ) + else: + num_tokens, hidden_size = q.shape + q = q.view(num_tokens, num_heads, head_size) + k = k.view(num_tokens, num_heads, head_size) + v = v.view(num_tokens, num_heads, head_size) + self_outputs = ops.ixinfer_flash_attn_unpad( + q, + k, + v, + cu_seq_lens, + cu_seq_lens, + max_seq_len, + max_seq_len, + is_causal=False, + atten_scale=1 / math.sqrt(head_size), + ) + self_outputs = (self_outputs,) + attention_output = self_outputs[0].view(num_tokens, hidden_size) + + attention_output = self.output(attention_output, hidden_states) + outputs = (attention_output,) + self_outputs[ + 1: + ] # add attentions if we output them + return outputs + + +class BertIntermediate(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] + else: + self.intermediate_act_fn = config.hidden_act + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + ori_shape = list(hidden_states.shape) + ori_shape[-1] = self.dense.weight.shape[0] + + hidden_states = ops.act_bias_mm( + mat1=hidden_states.view(-1, hidden_states.shape[-1]), + mat2=self.dense.weight, + bias=self.dense.bias, + act_type="gelu", + trans_format="TN", + ) + return hidden_states.view(*ori_shape) + + +class BertOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.eps = config.layer_norm_eps + + def forward( + self, hidden_states: torch.Tensor, input_tensor: torch.Tensor + ) -> torch.Tensor: + hidden_states = ops.linear(hidden_states, self.dense.weight, self.dense.bias) + hidden_states = ops.bert_add_norm( + input=hidden_states, + residual=input_tensor, + ln_weight=self.LayerNorm.weight, + ln_bias=self.LayerNorm.bias, + epsilon=self.eps, + ) + return hidden_states + + +class BertLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = BertAttention(config) + self.is_decoder = config.is_decoder + self.add_cross_attention = config.add_cross_attention + if self.add_cross_attention: + if not self.is_decoder: + raise ValueError( + f"{self} should be used as a decoder model if cross attention is added" + ) + self.crossattention = BertAttention( + config, position_embedding_type="absolute" + ) + self.intermediate = BertIntermediate(config) + self.output = BertOutput(config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + output_attentions: Optional[bool] = False, + cu_seq_lens: Optional[torch.Tensor] = None, + enable_unpad: Optional[bool] = False, + max_seq_len: Optional[int] = None, + ) -> Tuple[torch.Tensor]: + self_attention_outputs = self.attention( + hidden_states, + attention_mask, + cu_seq_lens=cu_seq_lens, + enable_unpad=enable_unpad, + max_seq_len=max_seq_len, + ) + attention_output = self_attention_outputs[0] + + outputs = self_attention_outputs[ + 1: + ] # add self attentions if we output attention weights + + layer_output = self.feed_forward_chunk(attention_output) + + outputs = (layer_output,) + outputs + + return outputs + + def feed_forward_chunk(self, attention_output): + intermediate_output = self.intermediate(attention_output) + layer_output = self.output(intermediate_output, attention_output) + return layer_output + + +class BertEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList( + [BertLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = False, + output_hidden_states: Optional[bool] = False, + return_dict: Optional[bool] = True, + cu_seq_lens: Optional[torch.Tensor] = None, + enable_unpad: Optional[bool] = False, + max_seq_len: Optional[int] = None, + ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]: + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + all_cross_attentions = ( + () if output_attentions and self.config.add_cross_attention else None + ) + + next_decoder_cache = () if use_cache else None + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_outputs = layer_module( + hidden_states, + attention_mask, + None, + None, + None, + None, + output_attentions, + cu_seq_lens=cu_seq_lens, + enable_unpad=enable_unpad, + max_seq_len=max_seq_len, + ) + + hidden_states = layer_outputs[0] + if use_cache: + next_decoder_cache += (layer_outputs[-1],) + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + if self.config.add_cross_attention: + all_cross_attentions = all_cross_attentions + (layer_outputs[2],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + next_decoder_cache, + all_hidden_states, + all_self_attentions, + all_cross_attentions, + ] + if v is not None + ) + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=next_decoder_cache, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + cross_attentions=all_cross_attentions, + ) + + +class BertPooler(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + first_token_tensor = hidden_states[:, 0] + pooled_output = self.dense(first_token_tensor) + pooled_output = self.activation(pooled_output) + return pooled_output + + +class BertPredictionHeadTransform(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + if isinstance(config.hidden_act, str): + self.transform_act_fn = ACT2FN[config.hidden_act] + else: + self.transform_act_fn = config.hidden_act + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.transform_act_fn(hidden_states) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + +class BertLMPredictionHead(nn.Module): + def __init__(self, config): + super().__init__() + self.transform = BertPredictionHeadTransform(config) + + # The output weights are the same as the input embeddings, but there is + # an output-only bias for each token. + self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + self.bias = nn.Parameter(torch.zeros(config.vocab_size)) + + # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings` + self.decoder.bias = self.bias + + def _tie_weights(self): + self.decoder.bias = self.bias + + def forward(self, hidden_states): + hidden_states = self.transform(hidden_states) + hidden_states = self.decoder(hidden_states) + return hidden_states + + +class BertOnlyMLMHead(nn.Module): + def __init__(self, config): + super().__init__() + self.predictions = BertLMPredictionHead(config) + + def forward(self, sequence_output: torch.Tensor) -> torch.Tensor: + prediction_scores = self.predictions(sequence_output) + return prediction_scores + + +class BertOnlyNSPHead(nn.Module): + def __init__(self, config): + super().__init__() + self.seq_relationship = nn.Linear(config.hidden_size, 2) + + def forward(self, pooled_output): + seq_relationship_score = self.seq_relationship(pooled_output) + return seq_relationship_score + + +class BertPreTrainingHeads(nn.Module): + def __init__(self, config): + super().__init__() + self.predictions = BertLMPredictionHead(config) + self.seq_relationship = nn.Linear(config.hidden_size, 2) + + def forward(self, sequence_output, pooled_output): + prediction_scores = self.predictions(sequence_output) + seq_relationship_score = self.seq_relationship(pooled_output) + return prediction_scores, seq_relationship_score + + +class BertPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = BertConfig + load_tf_weights = load_tf_weights_in_bert + base_model_prefix = "bert" + supports_gradient_checkpointing = True + _supports_sdpa = True + + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, nn.Linear): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + +@dataclass +class BertForPreTrainingOutput(ModelOutput): + """ + Output type of [`BertForPreTraining`]. + + Args: + loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): + Total loss as the sum of the masked language modeling loss and the next sequence prediction + (classification) loss. + prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`): + Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation + before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of + shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + prediction_logits: torch.FloatTensor = None + seq_relationship_logits: torch.FloatTensor = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + + +BERT_START_DOCSTRING = r""" + + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`BertConfig`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +BERT_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `({0})`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.FloatTensor` of shape `({0})`or `(batch_size, sequence_length, target_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*): + Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, + 1]`: + + - 0 corresponds to a *sentence A* token, + - 1 corresponds to a *sentence B* token. + + [What are token type IDs?](../glossary#token-type-ids) + position_ids (`torch.LongTensor` of shape `({0})`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.", + BERT_START_DOCSTRING, +) +class BertModel(BertPreTrainedModel): + """ + + The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of + cross-attention is added between the self-attention layers, following the architecture described in [Attention is + all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, + Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. + + To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set + to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and + `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass. + """ + + _no_split_modules = ["BertEmbeddings", "BertLayer"] + + def __init__(self, config, add_pooling_layer=True): + super().__init__(config) + self.config = config + + self.embeddings = BertEmbeddings(config) + self.encoder = BertEncoder(config) + + self.pooler = BertPooler(config) if add_pooling_layer else None + + self.attn_implementation = config._attn_implementation + self.position_embedding_type = config.position_embedding_type + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + def _prune_heads(self, heads_to_prune): + """ + Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base + class PreTrainedModel + """ + for layer, heads in heads_to_prune.items(): + self.encoder.layer[layer].attention.prune_heads(heads) + + @add_start_docstrings_to_model_forward( + BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=BaseModelOutputWithPoolingAndCrossAttentions, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cu_seq_lens: Optional[torch.Tensor] = None, + enable_unpad: Optional[bool] = False, + max_seq_len: Optional[int] = None, + ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + past_key_values_length=0, + ) + + encoder_outputs = self.encoder( + embedding_output, + attention_mask=attention_mask, + head_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_values=None, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cu_seq_lens=cu_seq_lens, + enable_unpad=enable_unpad, + max_seq_len=max_seq_len, + ) + + sequence_output = encoder_outputs[0] + if self.pooler is not None and enable_unpad: + raise NotImplementedError() + pooled_output = ( + self.pooler(sequence_output) if self.pooler is not None else None + ) + + if not return_dict: + return (sequence_output, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=encoder_outputs.past_key_values, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + cross_attentions=encoder_outputs.cross_attentions, + ) + + +@add_start_docstrings( + """ + Bert Model with two heads on top as done during the pretraining: a `masked language modeling` head and a `next + sentence prediction (classification)` head. + """, + BERT_START_DOCSTRING, +) +class BertForPreTraining(BertPreTrainedModel): + _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"] + + def __init__(self, config): + super().__init__(config) + + self.bert = BertModel(config) + self.cls = BertPreTrainingHeads(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.cls.predictions.decoder + + def set_output_embeddings(self, new_embeddings): + self.cls.predictions.decoder = new_embeddings + self.cls.predictions.bias = new_embeddings.bias + + @add_start_docstrings_to_model_forward( + BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @replace_return_docstrings( + output_type=BertForPreTrainingOutput, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + next_sentence_label: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor], BertForPreTrainingOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), + the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the next sequence prediction (classification) loss. Input should be a sequence + pair (see `input_ids` docstring) Indices should be in `[0, 1]`: + + - 0 indicates sequence B is a continuation of sequence A, + - 1 indicates sequence B is a random sequence. + kwargs (`Dict[str, any]`, *optional*, defaults to `{}`): + Used to hide legacy arguments that have been deprecated. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, BertForPreTraining + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased") + >>> model = BertForPreTraining.from_pretrained("google-bert/bert-base-uncased") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + >>> outputs = model(**inputs) + + >>> prediction_logits = outputs.prediction_logits + >>> seq_relationship_logits = outputs.seq_relationship_logits + ``` + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output, pooled_output = outputs[:2] + prediction_scores, seq_relationship_score = self.cls( + sequence_output, pooled_output + ) + + total_loss = None + if labels is not None and next_sentence_label is not None: + loss_fct = CrossEntropyLoss() + masked_lm_loss = loss_fct( + prediction_scores.view(-1, self.config.vocab_size), labels.view(-1) + ) + next_sentence_loss = loss_fct( + seq_relationship_score.view(-1, 2), next_sentence_label.view(-1) + ) + total_loss = masked_lm_loss + next_sentence_loss + + if not return_dict: + output = (prediction_scores, seq_relationship_score) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return BertForPreTrainingOutput( + loss=total_loss, + prediction_logits=prediction_scores, + seq_relationship_logits=seq_relationship_score, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """Bert Model with a `language modeling` head on top for CLM fine-tuning.""", + BERT_START_DOCSTRING, +) +class BertLMHeadModel(BertPreTrainedModel, GenerationMixin): + _tied_weights_keys = [ + "cls.predictions.decoder.bias", + "cls.predictions.decoder.weight", + ] + + def __init__(self, config): + super().__init__(config) + + if not config.is_decoder: + logger.warning( + "If you want to use `BertLMHeadModel` as a standalone, add `is_decoder=True.`" + ) + + self.bert = BertModel(config, add_pooling_layer=False) + self.cls = BertOnlyMLMHead(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.cls.predictions.decoder + + def set_output_embeddings(self, new_embeddings): + self.cls.predictions.decoder = new_embeddings + self.cls.predictions.bias = new_embeddings.bias + + @add_start_docstrings_to_model_forward( + BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=CausalLMOutputWithCrossAttentions, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.Tensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: + r""" + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if + the model is configured as a decoder. + encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in + the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in + `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are + ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]` + past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + if labels is not None: + use_cache = False + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + lm_loss = None + if labels is not None: + # we are doing next-token prediction; shift prediction scores and input ids by one + shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous() + labels = labels[:, 1:].contiguous() + loss_fct = CrossEntropyLoss() + lm_loss = loss_fct( + shifted_prediction_scores.view(-1, self.config.vocab_size), + labels.view(-1), + ) + + if not return_dict: + output = (prediction_scores,) + outputs[2:] + return ((lm_loss,) + output) if lm_loss is not None else output + + return CausalLMOutputWithCrossAttentions( + loss=lm_loss, + logits=prediction_scores, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + cross_attentions=outputs.cross_attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + use_cache=True, + **model_kwargs, + ): + input_shape = input_ids.shape + # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly + if attention_mask is None: + attention_mask = input_ids.new_ones(input_shape) + + # cut decoder_input_ids if past_key_values is used + if past_key_values is not None: + past_length = past_key_values[0][0].shape[2] + + # Some generation methods already pass only the last input ID + if input_ids.shape[1] > past_length: + remove_prefix_length = past_length + else: + # Default to old behavior: keep only final ID + remove_prefix_length = input_ids.shape[1] - 1 + + input_ids = input_ids[:, remove_prefix_length:] + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "use_cache": use_cache, + } + + def _reorder_cache(self, past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += ( + tuple( + past_state.index_select(0, beam_idx.to(past_state.device)) + for past_state in layer_past + ), + ) + return reordered_past + + +@add_start_docstrings( + """Bert Model with a `language modeling` head on top.""", BERT_START_DOCSTRING +) +class BertForMaskedLM(BertPreTrainedModel): + _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"] + + def __init__(self, config): + super().__init__(config) + + if config.is_decoder: + logger.warning( + "If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for " + "bi-directional self-attention." + ) + + self.bert = BertModel(config, add_pooling_layer=False) + self.cls = BertOnlyMLMHead(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.cls.predictions.decoder + + def set_output_embeddings(self, new_embeddings): + self.cls.predictions.decoder = new_embeddings + self.cls.predictions.bias = new_embeddings.bias + + @add_start_docstrings_to_model_forward( + BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=MaskedLMOutput, + config_class=_CONFIG_FOR_DOC, + expected_output="'paris'", + expected_loss=0.88, + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the + loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + """ + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() # -100 index = padding token + masked_lm_loss = loss_fct( + prediction_scores.view(-1, self.config.vocab_size), labels.view(-1) + ) + + if not return_dict: + output = (prediction_scores,) + outputs[2:] + return ( + ((masked_lm_loss,) + output) if masked_lm_loss is not None else output + ) + + return MaskedLMOutput( + loss=masked_lm_loss, + logits=prediction_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, input_ids, attention_mask=None, **model_kwargs + ): + input_shape = input_ids.shape + effective_batch_size = input_shape[0] + + # add a dummy token + if self.config.pad_token_id is None: + raise ValueError("The PAD token should be defined for generation") + + attention_mask = torch.cat( + [attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], + dim=-1, + ) + dummy_token = torch.full( + (effective_batch_size, 1), + self.config.pad_token_id, + dtype=torch.long, + device=input_ids.device, + ) + input_ids = torch.cat([input_ids, dummy_token], dim=1) + + return {"input_ids": input_ids, "attention_mask": attention_mask} + + +@add_start_docstrings( + """Bert Model with a `next sentence prediction (classification)` head on top.""", + BERT_START_DOCSTRING, +) +class BertForNextSentencePrediction(BertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.bert = BertModel(config) + self.cls = BertOnlyNSPHead(config) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward( + BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @replace_return_docstrings( + output_type=NextSentencePredictorOutput, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs, + ) -> Union[Tuple[torch.Tensor], NextSentencePredictorOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair + (see `input_ids` docstring). Indices should be in `[0, 1]`: + + - 0 indicates sequence B is a continuation of sequence A, + - 1 indicates sequence B is a random sequence. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, BertForNextSentencePrediction + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased") + >>> model = BertForNextSentencePrediction.from_pretrained("google-bert/bert-base-uncased") + + >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." + >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light." + >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt") + + >>> outputs = model(**encoding, labels=torch.LongTensor([1])) + >>> logits = outputs.logits + >>> assert logits[0, 0] < logits[0, 1] # next sentence was random + ``` + """ + + if "next_sentence_label" in kwargs: + warnings.warn( + "The `next_sentence_label` argument is deprecated and will be removed in a future version, use" + " `labels` instead.", + FutureWarning, + ) + labels = kwargs.pop("next_sentence_label") + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = outputs[1] + + seq_relationship_scores = self.cls(pooled_output) + + next_sentence_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + next_sentence_loss = loss_fct( + seq_relationship_scores.view(-1, 2), labels.view(-1) + ) + + if not return_dict: + output = (seq_relationship_scores,) + outputs[2:] + return ( + ((next_sentence_loss,) + output) + if next_sentence_loss is not None + else output + ) + + return NextSentencePredictorOutput( + loss=next_sentence_loss, + logits=seq_relationship_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled + output) e.g. for GLUE tasks. + """, + BERT_START_DOCSTRING, +) +class BertForSequenceClassification(BertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.config = config + + self.bert = BertModel(config) + classifier_dropout = ( + config.classifier_dropout + if config.classifier_dropout is not None + else config.hidden_dropout_prob + ) + self.dropout = nn.Dropout(classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward( + BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_SEQUENCE_CLASSIFICATION, + output_type=SequenceClassifierOutput, + config_class=_CONFIG_FOR_DOC, + expected_output=_SEQ_CLASS_EXPECTED_OUTPUT, + expected_loss=_SEQ_CLASS_EXPECTED_LOSS, + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = outputs[1] + + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + Bert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a + softmax) e.g. for RocStories/SWAG tasks. + """, + BERT_START_DOCSTRING, +) +class BertForMultipleChoice(BertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.bert = BertModel(config) + classifier_dropout = ( + config.classifier_dropout + if config.classifier_dropout is not None + else config.hidden_dropout_prob + ) + self.dropout = nn.Dropout(classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, 1) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward( + BERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=MultipleChoiceModelOutput, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor], MultipleChoiceModelOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., + num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See + `input_ids` above) + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + num_choices = ( + input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + ) + + input_ids = ( + input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None + ) + attention_mask = ( + attention_mask.view(-1, attention_mask.size(-1)) + if attention_mask is not None + else None + ) + token_type_ids = ( + token_type_ids.view(-1, token_type_ids.size(-1)) + if token_type_ids is not None + else None + ) + position_ids = ( + position_ids.view(-1, position_ids.size(-1)) + if position_ids is not None + else None + ) + inputs_embeds = ( + inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) + if inputs_embeds is not None + else None + ) + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = outputs[1] + + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + reshaped_logits = logits.view(-1, num_choices) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(reshaped_logits, labels) + + if not return_dict: + output = (reshaped_logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return MultipleChoiceModelOutput( + loss=loss, + logits=reshaped_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + Bert Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for + Named-Entity-Recognition (NER) tasks. + """, + BERT_START_DOCSTRING, +) +class BertForTokenClassification(BertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.bert = BertModel(config, add_pooling_layer=False) + classifier_dropout = ( + config.classifier_dropout + if config.classifier_dropout is not None + else config.hidden_dropout_prob + ) + self.dropout = nn.Dropout(classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward( + BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_TOKEN_CLASSIFICATION, + output_type=TokenClassifierOutput, + config_class=_CONFIG_FOR_DOC, + expected_output=_TOKEN_CLASS_EXPECTED_OUTPUT, + expected_loss=_TOKEN_CLASS_EXPECTED_LOSS, + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + Bert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear + layers on top of the hidden-states output to compute `span start logits` and `span end logits`). + """, + BERT_START_DOCSTRING, +) +class BertForQuestionAnswering(BertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.bert = BertModel(config, add_pooling_layer=False) + self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward( + BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + ) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_QA, + output_type=QuestionAnsweringModelOutput, + config_class=_CONFIG_FOR_DOC, + qa_target_start_index=_QA_TARGET_START_INDEX, + qa_target_end_index=_QA_TARGET_END_INDEX, + expected_output=_QA_EXPECTED_OUTPUT, + expected_loss=_QA_EXPECTED_LOSS, + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + start_positions: Optional[torch.Tensor] = None, + end_positions: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cu_seq_lens: Optional[torch.Tensor] = None, + max_seq_len: Optional[int] = None, + ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]: + r""" + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + if cu_seq_lens is not None: + enable_unpad = True + else: + enable_unpad = False + outputs = self.bert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cu_seq_lens=cu_seq_lens, + enable_unpad=enable_unpad, + max_seq_len=max_seq_len, + ) + + sequence_output = outputs[0] + + logits = ops.linear( + sequence_output, self.qa_outputs.weight, self.qa_outputs.bias + ) + + if enable_unpad: + start_logits, end_logits = ops.bert_unpack_start_end_logits( + logits, cu_seq_lens, max_seq_len + ) + else: + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + return QuestionAnsweringModelOutput( + loss=None, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=None, + attentions=None, + ) diff --git a/ixformer_sdk/contrib/transformers/models/t5/__init__.py b/ixformer_sdk/contrib/transformers/models/t5/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/transformers/models/t5/configuration_t5.py b/ixformer_sdk/contrib/transformers/models/t5/configuration_t5.py new file mode 100644 index 00000000..bb6d61b1 --- /dev/null +++ b/ixformer_sdk/contrib/transformers/models/t5/configuration_t5.py @@ -0,0 +1,174 @@ +# coding=utf-8 +# Copyright 2020, The T5 Authors and HuggingFace Inc. +# +# 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. +""" T5 model configuration""" +from typing import Mapping + +from transformers.configuration_utils import PretrainedConfig +from transformers.onnx import OnnxSeq2SeqConfigWithPast +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +T5_PRETRAINED_CONFIG_ARCHIVE_MAP = { + "t5-small": "https://huggingface.co/t5-small/resolve/main/config.json", + "t5-base": "https://huggingface.co/t5-base/resolve/main/config.json", + "t5-large": "https://huggingface.co/t5-large/resolve/main/config.json", + "t5-3b": "https://huggingface.co/t5-3b/resolve/main/config.json", + "t5-11b": "https://huggingface.co/t5-11b/resolve/main/config.json", +} + + +class T5Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`T5Model`] or a [`TFT5Model`]. It is used to + instantiate a T5 model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the T5 + [t5-small](https://huggingface.co/t5-small) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Arguments: + vocab_size (`int`, *optional*, defaults to 32128): + Vocabulary size of the T5 model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`T5Model`] or [`TFT5Model`]. + d_model (`int`, *optional*, defaults to 512): + Size of the encoder layers and the pooler layer. + d_kv (`int`, *optional*, defaults to 64): + Size of the key, query, value projections per attention head. The `inner_dim` of the projection layer will + be defined as `num_heads * d_kv`. + d_ff (`int`, *optional*, defaults to 2048): + Size of the intermediate feed forward layer in each `T5Block`. + num_layers (`int`, *optional*, defaults to 6): + Number of hidden layers in the Transformer encoder. + num_decoder_layers (`int`, *optional*): + Number of hidden layers in the Transformer decoder. Will use the same value as `num_layers` if not set. + num_heads (`int`, *optional*, defaults to 8): + Number of attention heads for each attention layer in the Transformer encoder. + relative_attention_num_buckets (`int`, *optional*, defaults to 32): + The number of buckets to use for each attention layer. + relative_attention_max_distance (`int`, *optional*, defaults to 128): + The maximum distance of the longer sequences for the bucket separation. + dropout_rate (`float`, *optional*, defaults to 0.1): + The ratio for all dropout layers. + layer_norm_eps (`float`, *optional*, defaults to 1e-6): + The epsilon used by the layer normalization layers. + initializer_factor (`float`, *optional*, defaults to 1): + A factor for initializing all weight matrices (should be kept to 1, used internally for initialization + testing). + feed_forward_proj (`string`, *optional*, defaults to `"relu"`): + Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. T5v1.1 uses the + `"gated-gelu"` feed forward projection. Original T5 uses `"relu"`. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). + """ + model_type = "t5" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "hidden_size": "d_model", + "num_attention_heads": "num_heads", + "num_hidden_layers": "num_layers", + } + + def __init__( + self, + vocab_size=32128, + d_model=512, + d_kv=64, + d_ff=2048, + num_layers=6, + num_decoder_layers=None, + num_heads=8, + relative_attention_num_buckets=32, + relative_attention_max_distance=128, + dropout_rate=0.1, + layer_norm_epsilon=1e-6, + initializer_factor=1.0, + feed_forward_proj="relu", + is_encoder_decoder=True, + use_cache=True, + pad_token_id=0, + eos_token_id=1, + **kwargs, + ): + self.vocab_size = vocab_size + self.d_model = d_model + self.d_kv = d_kv + self.d_ff = d_ff + self.num_layers = num_layers + self.num_decoder_layers = ( + num_decoder_layers if num_decoder_layers is not None else self.num_layers + ) # default = symmetry + self.num_heads = num_heads + self.relative_attention_num_buckets = relative_attention_num_buckets + self.relative_attention_max_distance = relative_attention_max_distance + self.dropout_rate = dropout_rate + self.layer_norm_epsilon = layer_norm_epsilon + self.initializer_factor = initializer_factor + self.feed_forward_proj = feed_forward_proj + self.use_cache = use_cache + + act_info = self.feed_forward_proj.split("-") + self.dense_act_fn = act_info[-1] + self.is_gated_act = act_info[0] == "gated" + + if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2: + raise ValueError( + f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer." + "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. " + "'gated-gelu' or 'relu'" + ) + + # for backwards compatibility + if feed_forward_proj == "gated-gelu": + self.dense_act_fn = "gelu_new" + + super().__init__( + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + is_encoder_decoder=is_encoder_decoder, + **kwargs, + ) + + +class T5OnnxConfig(OnnxSeq2SeqConfigWithPast): + @property + def inputs(self) -> Mapping[str, Mapping[int, str]]: + common_inputs = { + "input_ids": {0: "batch", 1: "encoder_sequence"}, + "attention_mask": {0: "batch", 1: "encoder_sequence"}, + } + if self.use_past: + common_inputs["attention_mask"][1] = "past_encoder_sequence + sequence" + common_inputs["decoder_input_ids"] = {0: "batch"} + common_inputs["decoder_attention_mask"] = { + 0: "batch", + 1: "past_decoder_sequence + sequence", + } + else: + common_inputs["decoder_input_ids"] = {0: "batch", 1: "decoder_sequence"} + common_inputs["decoder_attention_mask"] = { + 0: "batch", + 1: "decoder_sequence", + } + + if self.use_past: + self.fill_with_past_key_values_(common_inputs, direction="inputs") + + return common_inputs + + @property + def default_onnx_opset(self) -> int: + return 13 diff --git a/ixformer_sdk/contrib/transformers/models/t5/ixformer_modeling_t5.py b/ixformer_sdk/contrib/transformers/models/t5/ixformer_modeling_t5.py new file mode 100644 index 00000000..c8155325 --- /dev/null +++ b/ixformer_sdk/contrib/transformers/models/t5/ixformer_modeling_t5.py @@ -0,0 +1,315 @@ +import torch +from transformers.activations import NewGELUActivation + +import ixformer + + +def self_attention_forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + layer_head_mask=None, + past_key_value=None, + use_cache=False, + output_attentions=False, +): + assert output_attentions is False + assert layer_head_mask is None + + normed_hidden_states = self.layer_norm(hidden_states) + + if not hasattr(self, "qkv_weight"): + self.qkv_weight = torch.cat( + [ + self.SelfAttention.q.weight, + self.SelfAttention.k.weight, + self.SelfAttention.v.weight, + ], + dim=0, + ) + self.qkv_bias = None + + del self.SelfAttention.q.weight + del self.SelfAttention.k.weight + del self.SelfAttention.v.weight + + batch_size, seq_length = hidden_states.shape[:2] + real_seq_length = seq_length + if past_key_value is not None: + if len(past_key_value) != 2: + raise ValueError( + f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states" + ) + real_seq_length += past_key_value[0].shape[2] + key_length = real_seq_length + + def unshape(states): + """reshape""" + return ( + states.transpose(1, 2) + .contiguous() + .view(batch_size, -1, self.SelfAttention.inner_dim) + ) + + qkv = ixformer.functions.linear( + normed_hidden_states, self.qkv_weight, self.qkv_bias + ) + + if past_key_value is not None: + pask_key, past_value = past_key_value + ( + query_states, + key_states, + value_states, + ) = ixformer.functions.t5_split_qkv_update_kv_cache( + qkv, + pask_key, + past_value, + self.SelfAttention.n_heads, + self.SelfAttention.key_value_proj_dim, + ) + else: + query_states, key_states, value_states = ixformer.functions.t5_split_qkv( + qkv, self.SelfAttention.n_heads, self.SelfAttention.key_value_proj_dim + ) + + if position_bias is None: + if not self.SelfAttention.has_relative_attention_bias: + position_bias = torch.zeros( + (1, self.SelfAttention.n_heads, real_seq_length, key_length), + device=query_states.device, + dtype=query_states.dtype, + ) + else: + position_bias = self.SelfAttention.compute_bias( + real_seq_length, key_length, device=query_states.device + ) + + # if key and values are already calculated + # we want only the last query position bias + if past_key_value is not None: + position_bias = position_bias[:, :, -hidden_states.size(1) :, :] + + if attention_mask is not None: + # (batch_size, n_heads, seq_length, key_length) + position_bias = position_bias + attention_mask + + if self.SelfAttention.pruned_heads: + mask = torch.ones(position_bias.shape[1]) + mask[list(self.pruned_heads)] = 0 + position_bias_masked = position_bias[:, mask.bool()] + else: + position_bias_masked = position_bias + + attn_output = ixformer.functions.ixinfer_flash_attn_pad( + query_states.contiguous(), + key_states.contiguous(), + value_states.contiguous(), + mask=position_bias_masked.float().contiguous(), + atten_scale=1, + ) + attn_output = unshape(attn_output) + + attn_output = self.SelfAttention.o(attn_output) + + present_key_value_state = ( + (key_states, value_states) + if (self.SelfAttention.is_decoder and use_cache) + else None + ) + outputs = (attn_output,) + (present_key_value_state,) + (position_bias,) + + if output_attentions: + outputs = outputs + (None,) + hidden_states = attn_output + hidden_states + outputs = (hidden_states,) + outputs[1:] + + return outputs + + +def cross_attention_forward( + self, + hidden_states, + key_value_states, + attention_mask=None, + position_bias=None, + layer_head_mask=None, + past_key_value=None, + use_cache=False, + query_length=None, + output_attentions=False, +): + + assert output_attentions is False + assert layer_head_mask is None + + def unshape(states): + """reshape""" + return ( + states.transpose(1, 2) + .contiguous() + .view(batch_size, -1, self.EncDecAttention.inner_dim) + ) + + normed_hidden_states = self.layer_norm(hidden_states) + + # cross attn need key_value_states + assert key_value_states is not None + batch_size, seq_length = hidden_states.shape[:2] + real_seq_length = seq_length + + if past_key_value is not None: + if len(past_key_value) != 2: + raise ValueError( + f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states" + ) + real_seq_length += ( + past_key_value[0].shape[2] if query_length is None else query_length + ) + + key_length = ( + real_seq_length if key_value_states is None else key_value_states.shape[1] + ) + head_num, head_dim = ( + self.EncDecAttention.n_heads, + self.EncDecAttention.key_value_proj_dim, + ) + + query_states = ( + self.EncDecAttention.q(normed_hidden_states) + .view(batch_size, seq_length, head_num, head_dim) + .transpose(1, 2) + .contiguous() + ) + + if past_key_value is not None: + if past_key_value[0].shape[2] != key_value_states.shape[1]: + # checking that the `sequence_length` of the `past_key_value` is the same as + # the provided `key_value_states` to support prefix tuning + # cross-attn + # (batch_size, n_heads, seq_length, dim_per_head) + key_states = ( + self.EncDecAttention.k(key_value_states) + .view(batch_size, key_length, head_num, head_dim) + .transpose(1, 2) + ) + value_states = ( + self.EncDecAttention.v(key_value_states) + .view(batch_size, key_length, head_num, head_dim) + .transpose(1, 2) + ) + else: + # cross-attn + key_states = past_key_value[0] + value_states = past_key_value[1] + else: + key_states = ( + self.EncDecAttention.k(key_value_states) + .view(batch_size, key_length, head_num, head_dim) + .transpose(1, 2) + ) + value_states = ( + self.EncDecAttention.v(key_value_states) + .view(batch_size, key_length, head_num, head_dim) + .transpose(1, 2) + ) + + if not query_states.is_contiguous(): + query_states = query_states.contiguous() + + # TODO: fix this bug + if not value_states.is_contiguous(): + new_value_states = query_states.new_empty(value_states.shape) + new_value_states.copy_(value_states) + value_states = new_value_states + if not key_states.is_contiguous(): + numel = torch.numel(key_states) + new_key_states = query_states.new_empty([numel * 2])[:numel].view( + *list(key_states.shape) + ) + new_key_states.copy_(key_states) + key_states = new_key_states + + if position_bias is None: + if not self.EncDecAttention.has_relative_attention_bias: + position_bias = torch.zeros( + (1, self.EncDecAttention.n_heads, real_seq_length, key_length), + device=query_states.device, + dtype=query_states.dtype, + ) + else: + position_bias = self.EncDecAttention.compute_bias( + real_seq_length, key_length, device=query_states.device + ) + + # if key and values are already calculated + # we want only the last query position bias + if past_key_value is not None: + position_bias = position_bias[:, :, -hidden_states.size(1) :, :] + + if attention_mask is not None: + # (batch_size, n_heads, seq_length, key_length) + position_bias = position_bias + attention_mask + + if self.EncDecAttention.pruned_heads: + mask = torch.ones(position_bias.shape[1]) + mask[list(self.pruned_heads)] = 0 + position_bias_masked = position_bias[:, mask.bool()] + else: + position_bias_masked = position_bias + + attn_output = ixformer.functions.ixinfer_flash_attn_pad( + query_states, + key_states.contiguous(), + value_states.contiguous(), + mask=position_bias_masked.float().contiguous(), + atten_scale=1, + ) + attn_output = unshape(attn_output) + + attn_output = self.EncDecAttention.o(attn_output) + + present_key_value_state = ( + (key_states, value_states) + if (self.EncDecAttention.is_decoder and use_cache) + else None + ) + outputs = (attn_output,) + (present_key_value_state,) + (position_bias,) + + if output_attentions: + outputs = outputs + (None,) + hidden_states = attn_output + hidden_states + outputs = (hidden_states,) + outputs[1:] + + return outputs + + +def dense_gated_act_dense_forward(self, hidden_states): + if isinstance(self.act, NewGELUActivation): + if not hasattr(self, "wi"): + self.wi = torch.cat([self.wi_1.weight, self.wi_0.weight], dim=0) + del self.wi_1 + del self.wi_0 + hidden_states = ixformer.functions.linear(hidden_states, self.wi, None) + hidden_states = ixformer.functions.gelu_and_mul(hidden_states) + hidden_states = ixformer.functions.linear(hidden_states, self.wo.weight, None) + else: + hidden_gelu = self.act(self.wi_0(hidden_states)) + hidden_linear = self.wi_1(hidden_states) + hidden_states = hidden_gelu * hidden_linear + + hidden_states = self.dropout(hidden_states) + + # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32. + # See https://github.com/huggingface/transformers/issues/20287 + # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None`` + if ( + isinstance(self.wo.weight, torch.Tensor) + and hidden_states.dtype != self.wo.weight.dtype + and self.wo.weight.dtype != torch.int8 + ): + hidden_states = hidden_states.to(self.wo.weight.dtype) + + hidden_states = self.wo(hidden_states) + return hidden_states diff --git a/ixformer_sdk/contrib/transformers/models/t5/modeling_t5.py b/ixformer_sdk/contrib/transformers/models/t5/modeling_t5.py new file mode 100644 index 00000000..2c215bb5 --- /dev/null +++ b/ixformer_sdk/contrib/transformers/models/t5/modeling_t5.py @@ -0,0 +1,2644 @@ +# coding=utf-8 +# Copyright 2018 Mesh TensorFlow authors, T5 Authors and HuggingFace Inc. 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. +""" PyTorch T5 model.""" + + +import copy +import math +import os +import warnings +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from transformers.activations import ACT2FN +from transformers.modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPastAndCrossAttentions, + Seq2SeqLMOutput, + Seq2SeqModelOutput, + Seq2SeqQuestionAnsweringModelOutput, + Seq2SeqSequenceClassifierOutput, + TokenClassifierOutput, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.pytorch_utils import ( + ALL_LAYERNORM_LAYERS, + find_pruneable_heads_and_indices, + prune_linear_layer, +) +from transformers.utils import ( + DUMMY_INPUTS, + DUMMY_MASK, + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_torch_fx_proxy, + logging, + replace_return_docstrings, +) +from transformers.utils.model_parallel_utils import assert_device_map, get_device_map + +from .configuration_t5 import T5Config +from .ixformer_modeling_t5 import ( + cross_attention_forward, + dense_gated_act_dense_forward, + self_attention_forward, +) + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "T5Config" +_CHECKPOINT_FOR_DOC = "google-t5/t5-small" + +#################################################### +# This dict contains ids and associated url +# for the pretrained weights provided with the models +#################################################### +T5_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "google-t5/t5-small", + "google-t5/t5-base", + "google-t5/t5-large", + "google-t5/t5-3b", + "google-t5/t5-11b", + # See all T5 models at https://huggingface.co/models?filter=t5 +] + + +#################################################### +# This is a conversion method from TF 1.0 to PyTorch +# More details: https://medium.com/huggingface/from-tensorflow-to-pytorch-265f40ef2a28 +#################################################### +def load_tf_weights_in_t5(model, config, tf_checkpoint_path): + """Load tf checkpoints in a pytorch model.""" + try: + import re + + import numpy as np + import tensorflow as tf + except ImportError: + logger.error( + "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " + "https://www.tensorflow.org/install/ for installation instructions." + ) + raise + tf_path = os.path.abspath(tf_checkpoint_path) + logger.info(f"Converting TensorFlow checkpoint from {tf_path}") + # Load weights from TF model + init_vars = tf.train.list_variables(tf_path) + names = [] + tf_weights = {} + for name, shape in init_vars: + logger.info(f"Loading TF weight {name} with shape {shape}") + array = tf.train.load_variable(tf_path, name) + names.append(name) + tf_weights[name] = array + + for txt_name in names: + name = txt_name.split("/") + # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v + # which are not required for using pretrained model + if any( + n + in [ + "adam_v", + "adam_m", + "AdamWeightDecayOptimizer", + "AdamWeightDecayOptimizer_1", + "global_step", + ] + for n in name + ): + logger.info(f"Skipping {'/'.join(name)}") + tf_weights.pop(txt_name, None) + continue + if "_slot_" in name[-1]: + logger.info(f"Skipping {'/'.join(name)}") + tf_weights.pop(txt_name, None) + continue + pointer = model + array = tf_weights[txt_name] + + for m_name in name: + if re.fullmatch(r"[A-Za-z]+_\d+", m_name): + scope_names = re.split(r"_(\d+)", m_name) + else: + scope_names = [m_name] + if scope_names[0] in ["kernel", "scale", "embedding"]: + pointer = getattr(pointer, "weight") + elif scope_names[0] == "self_attention": + pointer = getattr(pointer, "layer") + pointer = pointer[0] + elif scope_names[0] == "enc_dec_attention": + pointer = getattr(pointer, "layer") + pointer = pointer[1] + elif scope_names[0] == "dense_relu_dense": + pointer = getattr(pointer, "layer") + pointer = pointer[2] + elif scope_names[0] == "rms_norm": + if hasattr(pointer, "layer_norm"): + pointer = getattr(pointer, "layer_norm") + elif hasattr(pointer, "final_layer_norm"): + pointer = getattr(pointer, "final_layer_norm") + elif scope_names[0] == "scale": + pointer = getattr(pointer, "weight") + elif scope_names[0] == "output_bias" or scope_names[0] == "beta": + pointer = getattr(pointer, "bias") + elif scope_names[0] == "squad": + pointer = getattr(pointer, "classifier") + elif scope_names[0] == "decoder" and name[1] == "logits": + continue + elif scope_names[0] == "logits": + pointer = getattr(pointer, "lm_head") + elif ( + scope_names[0] == "wi" + and len(scope_names) > 1 + and scope_names[1].isdigit() + ): + pointer = getattr(pointer, f"wi_{scope_names[1]}") + continue + else: + try: + pointer = getattr(pointer, scope_names[0]) + except AttributeError: + logger.info(f"Skipping {'/'.join(name)}") + continue + if len(scope_names) >= 2: + num = int(scope_names[1]) + pointer = pointer[num] + if scope_names[0] not in ["kernel", "scale", "embedding"]: + pointer = getattr(pointer, "weight") + if scope_names[0] != "embedding": + logger.info(f"Transposing numpy weight of shape {array.shape} for {name}") + array = np.transpose(array) + try: + if pointer.shape != array.shape: + raise ValueError( + f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched" + ) + except AssertionError as e: + e.args += (pointer.shape, array.shape) + raise + logger.info(f"Initialize PyTorch weight {name}") + pointer.data = torch.from_numpy(array.astype(np.float32)) + tf_weights.pop(txt_name, None) + + logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys())}.") + return model + + +#################################################### +# PyTorch Models are constructed by sub-classing +# - torch.nn.Module for the layers and +# - PreTrainedModel for the models (it-self a sub-class of nn.Module) +#################################################### +PARALLELIZE_DOCSTRING = r""" + This is an experimental feature and is a subject to change at a moment's notice. + + Uses a device map to distribute attention modules of the model across several devices. If no device map is given, + it will evenly distribute blocks across all devices. + + Args: + device_map (`Dict[int, list]`, optional, defaults to None): + A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always + automatically mapped to the first device (for esoteric reasons). That means that the first device should + have fewer attention modules mapped to it than other devices. For reference, the t5 models have the + following number of attention modules: + + - google-t5/t5-small: 6 + - google-t5/t5-base: 12 + - google-t5/t5-large: 24 + - google-t5/t5-3b: 24 + - google-t5/t5-11b: 24 + + Example: + + ```python + # Here is an example of a device map on a machine with 4 GPUs using google-t5/t5-3b, which has a total of 24 attention modules: + model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-3b") + device_map = { + 0: [0, 1, 2], + 1: [3, 4, 5, 6, 7, 8, 9], + 2: [10, 11, 12, 13, 14, 15, 16], + 3: [17, 18, 19, 20, 21, 22, 23], + } + model.parallelize(device_map) + ``` +""" +DEPARALLELIZE_DOCSTRING = r""" + Moves the model to cpu from a model parallel state. + + Example: + + ```python + # On a 4 GPU machine with google-t5/t5-3b: + model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-3b") + device_map = { + 0: [0, 1, 2], + 1: [3, 4, 5, 6, 7, 8, 9], + 2: [10, 11, 12, 13, 14, 15, 16], + 3: [17, 18, 19, 20, 21, 22, 23], + } + model.parallelize(device_map) # Splits the model across several devices + model.deparallelize() # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache() + ``` +""" + + +class T5LayerNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Construct a layernorm module in the T5 style. No bias and no subtraction of mean. + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + # T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean + # Square Layer Normalization https://arxiv.org/abs/1910.07467 thus varience is calculated + # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for + # half-precision inputs is done in fp32 + + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + + # convert into half-precision if necessary + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + + return self.weight * hidden_states + + +try: + from apex.normalization import FusedRMSNorm + + T5LayerNorm = FusedRMSNorm # noqa + + logger.info( + "Discovered apex.normalization.FusedRMSNorm - will use it instead of T5LayerNorm" + ) +except ImportError: + # using the normal T5LayerNorm + pass +except Exception: + logger.warning("discovered apex but it failed to load, falling back to T5LayerNorm") + pass + +ALL_LAYERNORM_LAYERS.append(T5LayerNorm) + + +class T5DenseActDense(nn.Module): + def __init__(self, config: T5Config): + super().__init__() + self.wi = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) + self.dropout = nn.Dropout(config.dropout_rate) + self.act = ACT2FN[config.dense_act_fn] + + def forward(self, hidden_states): + hidden_states = self.wi(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.dropout(hidden_states) + if ( + isinstance(self.wo.weight, torch.Tensor) + and hidden_states.dtype != self.wo.weight.dtype + and self.wo.weight.dtype != torch.int8 + ): + hidden_states = hidden_states.to(self.wo.weight.dtype) + hidden_states = self.wo(hidden_states) + return hidden_states + + +class T5DenseGatedActDense(nn.Module): + def __init__(self, config: T5Config): + super().__init__() + self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False) + self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) + self.dropout = nn.Dropout(config.dropout_rate) + self.act = ACT2FN[config.dense_act_fn] + + def forward(self, hidden_states): + return dense_gated_act_dense_forward(self, hidden_states) + hidden_gelu = self.act(self.wi_0(hidden_states)) + hidden_linear = self.wi_1(hidden_states) + hidden_states = hidden_gelu * hidden_linear + hidden_states = self.dropout(hidden_states) + + # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32. + # See https://github.com/huggingface/transformers/issues/20287 + # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None`` + if ( + isinstance(self.wo.weight, torch.Tensor) + and hidden_states.dtype != self.wo.weight.dtype + and self.wo.weight.dtype != torch.int8 + ): + hidden_states = hidden_states.to(self.wo.weight.dtype) + + hidden_states = self.wo(hidden_states) + return hidden_states + + +class T5LayerFF(nn.Module): + def __init__(self, config: T5Config): + super().__init__() + if config.is_gated_act: + self.DenseReluDense = T5DenseGatedActDense(config) + else: + self.DenseReluDense = T5DenseActDense(config) + + self.layer_norm = T5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward(self, hidden_states): + forwarded_states = self.layer_norm(hidden_states) + forwarded_states = self.DenseReluDense(forwarded_states) + hidden_states = hidden_states + self.dropout(forwarded_states) + return hidden_states + + +class T5Attention(nn.Module): + def __init__(self, config: T5Config, has_relative_attention_bias=False): + super().__init__() + self.is_decoder = config.is_decoder + self.has_relative_attention_bias = has_relative_attention_bias + self.relative_attention_num_buckets = config.relative_attention_num_buckets + self.relative_attention_max_distance = config.relative_attention_max_distance + self.d_model = config.d_model + self.key_value_proj_dim = config.d_kv + self.n_heads = config.num_heads + self.dropout = config.dropout_rate + self.inner_dim = self.n_heads * self.key_value_proj_dim + + # Mesh TensorFlow initialization to avoid scaling before softmax + self.q = nn.Linear(self.d_model, self.inner_dim, bias=False) + self.k = nn.Linear(self.d_model, self.inner_dim, bias=False) + self.v = nn.Linear(self.d_model, self.inner_dim, bias=False) + self.o = nn.Linear(self.inner_dim, self.d_model, bias=False) + + if self.has_relative_attention_bias: + self.relative_attention_bias = nn.Embedding( + self.relative_attention_num_buckets, self.n_heads + ) + self.pruned_heads = set() + self.gradient_checkpointing = False + + def prune_heads(self, heads): + if len(heads) == 0: + return + heads, index = find_pruneable_heads_and_indices( + heads, self.n_heads, self.key_value_proj_dim, self.pruned_heads + ) + # Prune linear layers + self.q = prune_linear_layer(self.q, index) + self.k = prune_linear_layer(self.k, index) + self.v = prune_linear_layer(self.v, index) + self.o = prune_linear_layer(self.o, index, dim=1) + # Update hyper params + self.n_heads = self.n_heads - len(heads) + self.inner_dim = self.key_value_proj_dim * self.n_heads + self.pruned_heads = self.pruned_heads.union(heads) + + @staticmethod + def _relative_position_bucket( + relative_position, bidirectional=True, num_buckets=32, max_distance=128 + ): + """ + Adapted from Mesh Tensorflow: + https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 + + Translate relative position to a bucket number for relative attention. The relative position is defined as + memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to + position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for + small absolute relative_position and larger buckets for larger absolute relative_positions. All relative + positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. + This should allow for more graceful generalization to longer sequences than the model has been trained on + + Args: + relative_position: an int32 Tensor + bidirectional: a boolean - whether the attention is bidirectional + num_buckets: an integer + max_distance: an integer + + Returns: + a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) + """ + relative_buckets = 0 + if bidirectional: + num_buckets //= 2 + relative_buckets += (relative_position > 0).to(torch.long) * num_buckets + relative_position = torch.abs(relative_position) + else: + relative_position = -torch.min( + relative_position, torch.zeros_like(relative_position) + ) + # now relative_position is in the range [0, inf) + + # half of the buckets are for exact increments in positions + max_exact = num_buckets // 2 + is_small = relative_position < max_exact + + # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance + relative_position_if_large = max_exact + ( + torch.log(relative_position.float() / max_exact) + / math.log(max_distance / max_exact) + * (num_buckets - max_exact) + ).to(torch.long) + relative_position_if_large = torch.min( + relative_position_if_large, + torch.full_like(relative_position_if_large, num_buckets - 1), + ) + + relative_buckets += torch.where( + is_small, relative_position, relative_position_if_large + ) + return relative_buckets + + def compute_bias(self, query_length, key_length, device=None): + """Compute binned relative position bias""" + if device is None: + device = self.relative_attention_bias.weight.device + context_position = torch.arange(query_length, dtype=torch.long, device=device)[ + :, None + ] + memory_position = torch.arange(key_length, dtype=torch.long, device=device)[ + None, : + ] + relative_position = ( + memory_position - context_position + ) # shape (query_length, key_length) + relative_position_bucket = self._relative_position_bucket( + relative_position, # shape (query_length, key_length) + bidirectional=(not self.is_decoder), + num_buckets=self.relative_attention_num_buckets, + max_distance=self.relative_attention_max_distance, + ) + values = self.relative_attention_bias( + relative_position_bucket + ) # shape (query_length, key_length, num_heads) + values = values.permute([2, 0, 1]).unsqueeze( + 0 + ) # shape (1, num_heads, query_length, key_length) + return values + + def forward( + self, + hidden_states, + mask=None, + key_value_states=None, + position_bias=None, + past_key_value=None, + layer_head_mask=None, + query_length=None, + use_cache=False, + output_attentions=False, + ): + """ + Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). + """ + # Input is (batch_size, seq_length, dim) + # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length) + # past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head) + batch_size, seq_length = hidden_states.shape[:2] + + real_seq_length = seq_length + + if past_key_value is not None: + if len(past_key_value) != 2: + raise ValueError( + f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states" + ) + real_seq_length += ( + past_key_value[0].shape[2] if query_length is None else query_length + ) + + key_length = ( + real_seq_length if key_value_states is None else key_value_states.shape[1] + ) + + def shape(states): + """projection""" + return states.view( + batch_size, -1, self.n_heads, self.key_value_proj_dim + ).transpose(1, 2) + + def unshape(states): + """reshape""" + return ( + states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim) + ) + + def project(hidden_states, proj_layer, key_value_states, past_key_value): + """projects hidden states correctly to key/query states""" + if key_value_states is None: + # self-attn + # (batch_size, n_heads, seq_length, dim_per_head) + hidden_states = shape(proj_layer(hidden_states)) + elif past_key_value is None: + # cross-attn + # (batch_size, n_heads, seq_length, dim_per_head) + hidden_states = shape(proj_layer(key_value_states)) + + if past_key_value is not None: + if key_value_states is None: + # self-attn + # (batch_size, n_heads, key_length, dim_per_head) + hidden_states = torch.cat([past_key_value, hidden_states], dim=2) + elif past_key_value.shape[2] != key_value_states.shape[1]: + # checking that the `sequence_length` of the `past_key_value` is the same as + # the provided `key_value_states` to support prefix tuning + # cross-attn + # (batch_size, n_heads, seq_length, dim_per_head) + hidden_states = shape(proj_layer(key_value_states)) + else: + # cross-attn + hidden_states = past_key_value + return hidden_states + + # get query states + query_states = shape( + self.q(hidden_states) + ) # (batch_size, n_heads, seq_length, dim_per_head) + + # get key/value states + key_states = project( + hidden_states, + self.k, + key_value_states, + past_key_value[0] if past_key_value is not None else None, + ) + value_states = project( + hidden_states, + self.v, + key_value_states, + past_key_value[1] if past_key_value is not None else None, + ) + + # compute scores + scores = torch.matmul( + query_states, key_states.transpose(3, 2) + ) # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 + + if position_bias is None: + if not self.has_relative_attention_bias: + position_bias = torch.zeros( + (1, self.n_heads, real_seq_length, key_length), + device=scores.device, + dtype=scores.dtype, + ) + if self.gradient_checkpointing and self.training: + position_bias.requires_grad = True + else: + position_bias = self.compute_bias( + real_seq_length, key_length, device=scores.device + ) + + # if key and values are already calculated + # we want only the last query position bias + if past_key_value is not None: + position_bias = position_bias[:, :, -hidden_states.size(1) :, :] + + if mask is not None: + position_bias = ( + position_bias + mask + ) # (batch_size, n_heads, seq_length, key_length) + + if self.pruned_heads: + mask = torch.ones(position_bias.shape[1]) + mask[list(self.pruned_heads)] = 0 + position_bias_masked = position_bias[:, mask.bool()] + else: + position_bias_masked = position_bias + + scores += position_bias_masked + attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as( + scores + ) # (batch_size, n_heads, seq_length, key_length) + attn_weights = nn.functional.dropout( + attn_weights, p=self.dropout, training=self.training + ) # (batch_size, n_heads, seq_length, key_length) + + # Mask heads if we want to + if layer_head_mask is not None: + attn_weights = attn_weights * layer_head_mask + + attn_output = unshape( + torch.matmul(attn_weights, value_states) + ) # (batch_size, seq_length, dim) + attn_output = self.o(attn_output) + + present_key_value_state = ( + (key_states, value_states) if (self.is_decoder and use_cache) else None + ) + outputs = (attn_output,) + (present_key_value_state,) + (position_bias,) + + if output_attentions: + outputs = outputs + (attn_weights,) + return outputs + + +class T5LayerSelfAttention(nn.Module): + def __init__(self, config, has_relative_attention_bias=False): + super().__init__() + self.SelfAttention = T5Attention( + config, has_relative_attention_bias=has_relative_attention_bias + ) + self.layer_norm = T5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + layer_head_mask=None, + past_key_value=None, + use_cache=False, + output_attentions=False, + ): + return self_attention_forward( + self, + hidden_states=hidden_states, + attention_mask=attention_mask, + position_bias=position_bias, + layer_head_mask=layer_head_mask, + past_key_value=past_key_value, + use_cache=use_cache, + output_attentions=output_attentions, + ) + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.SelfAttention( + normed_hidden_states, + mask=attention_mask, + position_bias=position_bias, + layer_head_mask=layer_head_mask, + past_key_value=past_key_value, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states = hidden_states + self.dropout(attention_output[0]) + outputs = (hidden_states,) + attention_output[ + 1: + ] # add attentions if we output them + return outputs + + +class T5LayerCrossAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.EncDecAttention = T5Attention(config, has_relative_attention_bias=False) + self.layer_norm = T5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) + self.dropout = nn.Dropout(config.dropout_rate) + + def forward( + self, + hidden_states, + key_value_states, + attention_mask=None, + position_bias=None, + layer_head_mask=None, + past_key_value=None, + use_cache=False, + query_length=None, + output_attentions=False, + ): + return cross_attention_forward( + self, + hidden_states=hidden_states, + key_value_states=key_value_states, + attention_mask=attention_mask, + position_bias=position_bias, + layer_head_mask=layer_head_mask, + past_key_value=past_key_value, + use_cache=use_cache, + query_length=query_length, + output_attentions=output_attentions, + ) + normed_hidden_states = self.layer_norm(hidden_states) + attention_output = self.EncDecAttention( + normed_hidden_states, + mask=attention_mask, + key_value_states=key_value_states, + position_bias=position_bias, + layer_head_mask=layer_head_mask, + past_key_value=past_key_value, + use_cache=use_cache, + query_length=query_length, + output_attentions=output_attentions, + ) + layer_output = hidden_states + self.dropout(attention_output[0]) + outputs = (layer_output,) + attention_output[ + 1: + ] # add attentions if we output them + return outputs + + +class T5Block(nn.Module): + def __init__(self, config, has_relative_attention_bias=False): + super().__init__() + self.is_decoder = config.is_decoder + self.layer = nn.ModuleList() + self.layer.append( + T5LayerSelfAttention( + config, has_relative_attention_bias=has_relative_attention_bias + ) + ) + if self.is_decoder: + self.layer.append(T5LayerCrossAttention(config)) + + self.layer.append(T5LayerFF(config)) + + def forward( + self, + hidden_states, + attention_mask=None, + position_bias=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + encoder_decoder_position_bias=None, + layer_head_mask=None, + cross_attn_layer_head_mask=None, + past_key_value=None, + use_cache=False, + output_attentions=False, + return_dict=True, + ): + if past_key_value is not None: + if not self.is_decoder: + logger.warning( + "`past_key_values` is passed to the encoder. Please make sure this is intended." + ) + expected_num_past_key_values = 2 if encoder_hidden_states is None else 4 + + if len(past_key_value) != expected_num_past_key_values: + raise ValueError( + f"There should be {expected_num_past_key_values} past states. " + f"{'2 (past / key) for cross attention. ' if expected_num_past_key_values == 4 else ''}" + f"Got {len(past_key_value)} past key / value states" + ) + + self_attn_past_key_value = past_key_value[:2] + cross_attn_past_key_value = past_key_value[2:] + else: + self_attn_past_key_value, cross_attn_past_key_value = None, None + + self_attention_outputs = self.layer[0]( + hidden_states, + attention_mask=attention_mask, + position_bias=position_bias, + layer_head_mask=layer_head_mask, + past_key_value=self_attn_past_key_value, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states, present_key_value_state = self_attention_outputs[:2] + attention_outputs = self_attention_outputs[ + 2: + ] # Keep self-attention outputs and relative position weights + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) + hidden_states = torch.clamp( + hidden_states, min=-clamp_value, max=clamp_value + ) + + do_cross_attention = self.is_decoder and encoder_hidden_states is not None + if do_cross_attention: + # the actual query length is unknown for cross attention + # if using past key value states. Need to inject it here + if present_key_value_state is not None: + query_length = present_key_value_state[0].shape[2] + else: + query_length = None + + cross_attention_outputs = self.layer[1]( + hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + position_bias=encoder_decoder_position_bias, + layer_head_mask=cross_attn_layer_head_mask, + past_key_value=cross_attn_past_key_value, + query_length=query_length, + use_cache=use_cache, + output_attentions=output_attentions, + ) + hidden_states = cross_attention_outputs[0] + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) + hidden_states = torch.clamp( + hidden_states, min=-clamp_value, max=clamp_value + ) + + # Combine self attn and cross attn key value states + if present_key_value_state is not None: + present_key_value_state = ( + present_key_value_state + cross_attention_outputs[1] + ) + + # Keep cross-attention outputs and relative position weights + attention_outputs = attention_outputs + cross_attention_outputs[2:] + + # Apply Feed Forward layer + hidden_states = self.layer[-1](hidden_states) + + # clamp inf values to enable fp16 training + if hidden_states.dtype == torch.float16: + clamp_value = torch.where( + torch.isinf(hidden_states).any(), + torch.finfo(hidden_states.dtype).max - 1000, + torch.finfo(hidden_states.dtype).max, + ) + hidden_states = torch.clamp( + hidden_states, min=-clamp_value, max=clamp_value + ) + + outputs = (hidden_states,) + + if use_cache: + outputs = outputs + (present_key_value_state,) + attention_outputs + else: + outputs = outputs + attention_outputs + + return outputs # hidden-states, present_key_value_states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) + + +class T5ClassificationHead(nn.Module): + """Head for sentence-level classification tasks.""" + + def __init__(self, config: T5Config): + super().__init__() + self.dense = nn.Linear(config.d_model, config.d_model) + self.dropout = nn.Dropout(p=config.classifier_dropout) + self.out_proj = nn.Linear(config.d_model, config.num_labels) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dropout(hidden_states) + hidden_states = self.dense(hidden_states) + hidden_states = torch.tanh(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.out_proj(hidden_states) + return hidden_states + + +class T5PreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = T5Config + load_tf_weights = load_tf_weights_in_t5 + base_model_prefix = "transformer" + is_parallelizable = True + supports_gradient_checkpointing = True + _no_split_modules = ["T5Block"] + _keep_in_fp32_modules = ["wo"] + + @property + def dummy_inputs(self): + input_ids = torch.tensor(DUMMY_INPUTS) + input_mask = torch.tensor(DUMMY_MASK) + dummy_inputs = { + "decoder_input_ids": input_ids, + "input_ids": input_ids, + "decoder_attention_mask": input_mask, + } + return dummy_inputs + + def _init_weights(self, module): + """Initialize the weights""" + factor = ( + self.config.initializer_factor + ) # Used for testing weights initialization + if isinstance(module, T5LayerNorm): + module.weight.data.fill_(factor * 1.0) + elif isinstance( + module, + ( + T5Model, + T5ForConditionalGeneration, + T5EncoderModel, + T5ForQuestionAnswering, + ), + ): + # Mesh TensorFlow embeddings initialization + # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624 + module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0) + if hasattr(module, "lm_head") and not self.config.tie_word_embeddings: + module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0) + if hasattr(module, "qa_outputs"): + module.qa_outputs.weight.data.normal_( + mean=0.0, std=factor * ((self.config.d_model) ** -0.5) + ) + module.qa_outputs.bias.data.zero_() + elif isinstance(module, T5ForTokenClassification): + if hasattr(module, "classifier"): + module.classifier.weight.data.normal_(mean=0.0, std=factor * 1.0) + module.classifier.bias.data.zero_() + elif isinstance(module, T5ClassificationHead): + module.dense.weight.data.normal_( + mean=0.0, std=factor * ((self.config.d_model) ** -0.5) + ) + if hasattr(module.dense, "bias") and module.dense.bias is not None: + module.dense.bias.data.zero_() + module.out_proj.weight.data.normal_( + mean=0.0, std=factor * ((self.config.d_model) ** -0.5) + ) + if hasattr(module.out_proj, "bias") and module.out_proj.bias is not None: + module.out_proj.bias.data.zero_() + elif isinstance(module, T5DenseActDense): + # Mesh TensorFlow FF initialization + # See https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L56 + # and https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L89 + module.wi.weight.data.normal_( + mean=0.0, std=factor * ((self.config.d_model) ** -0.5) + ) + if hasattr(module.wi, "bias") and module.wi.bias is not None: + module.wi.bias.data.zero_() + module.wo.weight.data.normal_( + mean=0.0, std=factor * ((self.config.d_ff) ** -0.5) + ) + if hasattr(module.wo, "bias") and module.wo.bias is not None: + module.wo.bias.data.zero_() + elif isinstance(module, T5DenseGatedActDense): + module.wi_0.weight.data.normal_( + mean=0.0, std=factor * ((self.config.d_model) ** -0.5) + ) + if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None: + module.wi_0.bias.data.zero_() + module.wi_1.weight.data.normal_( + mean=0.0, std=factor * ((self.config.d_model) ** -0.5) + ) + if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None: + module.wi_1.bias.data.zero_() + module.wo.weight.data.normal_( + mean=0.0, std=factor * ((self.config.d_ff) ** -0.5) + ) + if hasattr(module.wo, "bias") and module.wo.bias is not None: + module.wo.bias.data.zero_() + elif isinstance(module, T5Attention): + # Mesh TensorFlow attention initialization to avoid scaling before softmax + # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136 + d_model = self.config.d_model + key_value_proj_dim = self.config.d_kv + n_heads = self.config.num_heads + module.q.weight.data.normal_( + mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5) + ) + module.k.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) + module.v.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) + module.o.weight.data.normal_( + mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5) + ) + if module.has_relative_attention_bias: + module.relative_attention_bias.weight.data.normal_( + mean=0.0, std=factor * ((d_model) ** -0.5) + ) + + def _shift_right(self, input_ids): + decoder_start_token_id = self.config.decoder_start_token_id + pad_token_id = self.config.pad_token_id + + if decoder_start_token_id is None: + raise ValueError( + "self.model.config.decoder_start_token_id has to be defined. In T5 it is usually set to the pad_token_id. " + "See T5 docs for more information." + ) + + # shift inputs to the right + if is_torch_fx_proxy(input_ids): + # Item assignment is not supported natively for proxies. + shifted_input_ids = torch.full( + input_ids.shape[:-1] + (1,), decoder_start_token_id + ) + shifted_input_ids = torch.cat( + [shifted_input_ids, input_ids[..., :-1]], dim=-1 + ) + else: + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = decoder_start_token_id + + if pad_token_id is None: + raise ValueError("self.model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) + + return shifted_input_ids + + +class T5Stack(T5PreTrainedModel): + def __init__(self, config, embed_tokens=None): + super().__init__(config) + + self.embed_tokens = embed_tokens + self.is_decoder = config.is_decoder + + self.block = nn.ModuleList( + [ + T5Block(config, has_relative_attention_bias=bool(i == 0)) + for i in range(config.num_layers) + ] + ) + self.final_layer_norm = T5LayerNorm( + config.d_model, eps=config.layer_norm_epsilon + ) + self.dropout = nn.Dropout(config.dropout_rate) + + # Initialize weights and apply final processing + self.post_init() + # Model parallel + self.model_parallel = False + self.device_map = None + self.gradient_checkpointing = False + + @add_start_docstrings(PARALLELIZE_DOCSTRING) + def parallelize(self, device_map=None): + warnings.warn( + "`T5Stack.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model" + " with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" + " `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0," + " 'block.1': 1, ...}", + FutureWarning, + ) + # Check validity of device_map + self.device_map = ( + get_device_map(len(self.block), range(torch.cuda.device_count())) + if device_map is None + else device_map + ) + assert_device_map(self.device_map, len(self.block)) + self.model_parallel = True + self.first_device = ( + "cpu" + if "cpu" in self.device_map.keys() + else "cuda:" + str(min(self.device_map.keys())) + ) + self.last_device = "cuda:" + str(max(self.device_map.keys())) + # Load onto devices + for k, v in self.device_map.items(): + for layer in v: + cuda_device = "cuda:" + str(k) + self.block[layer] = self.block[layer].to(cuda_device) + + # Set embed_tokens to first layer + self.embed_tokens = self.embed_tokens.to(self.first_device) + # Set final layer norm to last device + self.final_layer_norm = self.final_layer_norm.to(self.last_device) + + @add_start_docstrings(DEPARALLELIZE_DOCSTRING) + def deparallelize(self): + warnings.warn( + "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", + FutureWarning, + ) + self.model_parallel = False + self.device_map = None + self.first_device = "cpu" + self.last_device = "cpu" + for i in range(len(self.block)): + self.block[i] = self.block[i].to("cpu") + self.embed_tokens = self.embed_tokens.to("cpu") + self.final_layer_norm = self.final_layer_norm.to("cpu") + torch.cuda.empty_cache() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, new_embeddings): + self.embed_tokens = new_embeddings + + def forward( + self, + input_ids=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + inputs_embeds=None, + head_mask=None, + cross_attn_head_mask=None, + past_key_values=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + # Model parallel + if self.model_parallel: + torch.cuda.set_device(self.first_device) + self.embed_tokens = self.embed_tokens.to(self.first_device) + use_cache = use_cache if use_cache is not None else self.config.use_cache + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if input_ids is not None and inputs_embeds is not None: + err_msg_prefix = "decoder_" if self.is_decoder else "" + raise ValueError( + f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}inputs_embeds at the same time" + ) + elif input_ids is not None: + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + err_msg_prefix = "decoder_" if self.is_decoder else "" + raise ValueError( + f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds" + ) + + if inputs_embeds is None: + if self.embed_tokens is None: + raise ValueError( + "You have to initialize the model with valid token embeddings" + ) + inputs_embeds = self.embed_tokens(input_ids) + + batch_size, seq_length = input_shape + + # required mask seq length can be calculated via length of past + mask_seq_length = ( + past_key_values[0][0].shape[2] + seq_length + if past_key_values is not None + else seq_length + ) + + if use_cache is True: + if not self.is_decoder: + raise ValueError( + f"`use_cache` can only be set to `True` if {self} is used as a decoder" + ) + + # initialize past_key_values with `None` if past does not exist + if past_key_values is None: + past_key_values = [None] * len(self.block) + + if attention_mask is None: + attention_mask = torch.ones( + batch_size, mask_seq_length, device=inputs_embeds.device + ) + + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + extended_attention_mask = self.get_extended_attention_mask( + attention_mask, input_shape + ) + + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.is_decoder and encoder_hidden_states is not None: + ( + encoder_batch_size, + encoder_sequence_length, + _, + ) = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones( + encoder_hidden_shape, device=inputs_embeds.device, dtype=torch.long + ) + encoder_extended_attention_mask = self.invert_attention_mask( + encoder_attention_mask + ) + else: + encoder_extended_attention_mask = None + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # Prepare head mask if needed + head_mask = self.get_head_mask(head_mask, self.config.num_layers) + cross_attn_head_mask = self.get_head_mask( + cross_attn_head_mask, self.config.num_layers + ) + present_key_value_states = () if use_cache else None + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + all_cross_attentions = () if (output_attentions and self.is_decoder) else None + position_bias = None + encoder_decoder_position_bias = None + + hidden_states = self.dropout(inputs_embeds) + + for i, (layer_module, past_key_value) in enumerate( + zip(self.block, past_key_values) + ): + layer_head_mask = head_mask[i] + cross_attn_layer_head_mask = cross_attn_head_mask[i] + # Model parallel + if self.model_parallel: + torch.cuda.set_device(hidden_states.device) + # Ensure that attention_mask is always on the same device as hidden_states + if attention_mask is not None: + attention_mask = attention_mask.to(hidden_states.device) + if position_bias is not None: + position_bias = position_bias.to(hidden_states.device) + if encoder_hidden_states is not None: + encoder_hidden_states = encoder_hidden_states.to( + hidden_states.device + ) + if encoder_extended_attention_mask is not None: + encoder_extended_attention_mask = ( + encoder_extended_attention_mask.to(hidden_states.device) + ) + if encoder_decoder_position_bias is not None: + encoder_decoder_position_bias = encoder_decoder_position_bias.to( + hidden_states.device + ) + if layer_head_mask is not None: + layer_head_mask = layer_head_mask.to(hidden_states.device) + if cross_attn_layer_head_mask is not None: + cross_attn_layer_head_mask = cross_attn_layer_head_mask.to( + hidden_states.device + ) + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + layer_module.forward, + hidden_states, + extended_attention_mask, + position_bias, + encoder_hidden_states, + encoder_extended_attention_mask, + encoder_decoder_position_bias, + layer_head_mask, + cross_attn_layer_head_mask, + None, # past_key_value is always None with gradient checkpointing + use_cache, + output_attentions, + ) + else: + layer_outputs = layer_module( + hidden_states, + attention_mask=extended_attention_mask, + position_bias=position_bias, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_extended_attention_mask, + encoder_decoder_position_bias=encoder_decoder_position_bias, + layer_head_mask=layer_head_mask, + cross_attn_layer_head_mask=cross_attn_layer_head_mask, + past_key_value=past_key_value, + use_cache=use_cache, + output_attentions=output_attentions, + ) + + # layer_outputs is a tuple with: + # hidden-states, key-value-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) + if use_cache is False: + layer_outputs = layer_outputs[:1] + (None,) + layer_outputs[1:] + + hidden_states, present_key_value_state = layer_outputs[:2] + + # We share the position biases between the layers - the first layer store them + # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights), + # (cross-attention position bias), (cross-attention weights) + position_bias = layer_outputs[2] + if self.is_decoder and encoder_hidden_states is not None: + encoder_decoder_position_bias = layer_outputs[ + 4 if output_attentions else 3 + ] + # append next layer key value states + if use_cache: + present_key_value_states = present_key_value_states + ( + present_key_value_state, + ) + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[3],) + if self.is_decoder: + all_cross_attentions = all_cross_attentions + (layer_outputs[5],) + + # Model Parallel: If it's the last layer for that device, put things on the next device + if self.model_parallel: + for k, v in self.device_map.items(): + if i == v[-1] and "cuda:" + str(k) != self.last_device: + hidden_states = hidden_states.to("cuda:" + str(k + 1)) + + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.dropout(hidden_states) + + # Add last layer + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + present_key_value_states, + all_hidden_states, + all_attentions, + all_cross_attentions, + ] + if v is not None + ) + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=present_key_value_states, + hidden_states=all_hidden_states, + attentions=all_attentions, + cross_attentions=all_cross_attentions, + ) + + +T5_START_DOCSTRING = r""" + + The T5 model was proposed in [Exploring the Limits of Transfer Learning with a Unified Text-to-Text + Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan + Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu. It's an encoder decoder transformer pre-trained in a + text-to-text denoising generative setting. + + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`T5Config`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +T5_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you + should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + [What are input IDs?](../glossary#input-ids) + + To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). + attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Indices of decoder input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are decoder input IDs?](../glossary#decoder-input-ids) + + T5 uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` + is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). + + To know more on how to prepare `decoder_input_ids` for pretraining take a look at [T5 + Training](./t5#training). + decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): + Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also + be used by default. + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules in the encoder. Mask values selected in `[0, + 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, + 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in + `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): + Tuple consists of (`last_hidden_state`, `optional`: *hidden_states*, `optional`: *attentions*) + `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` is a sequence of hidden states at + the output of the last layer of the encoder. Used in the cross-attention of the decoder. + past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded + representation. If `past_key_values` is used, optionally only the last `decoder_inputs_embeds` have to be + input (see `past_key_values`). This is useful if you want more control over how to convert + `decoder_input_ids` indices into associated vectors than the model's internal embedding lookup matrix. + + If `decoder_input_ids` and `decoder_inputs_embeds` are both unset, `decoder_inputs_embeds` takes the value + of `inputs_embeds`. + + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + +T5_ENCODER_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you + should be able to pad the inputs on both the right and the left. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for detail. + + To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). + attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + +# Warning message for FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask +__HEAD_MASK_WARNING_MSG = """ +The input argument `head_mask` was split into two arguments `head_mask` and `decoder_head_mask`. Currently, +`decoder_head_mask` is set to copy `head_mask`, but this feature is deprecated and will be removed in future versions. +If you do not want to use any `decoder_head_mask` now, please set `decoder_head_mask = torch.ones(num_layers, +num_heads)`. +""" + + +@add_start_docstrings( + "The bare T5 Model transformer outputting raw hidden-states without any specific head on top.", + T5_START_DOCSTRING, +) +class T5Model(T5PreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + "decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight", + ] + _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] + + def __init__(self, config: T5Config): + super().__init__(config) + self.shared = nn.Embedding(config.vocab_size, config.d_model) + + encoder_config = copy.deepcopy(config) + encoder_config.is_decoder = False + encoder_config.use_cache = False + encoder_config.is_encoder_decoder = False + self.encoder = T5Stack(encoder_config, self.shared) + + decoder_config = copy.deepcopy(config) + decoder_config.is_decoder = True + decoder_config.is_encoder_decoder = False + decoder_config.num_layers = config.num_decoder_layers + self.decoder = T5Stack(decoder_config, self.shared) + + # Initialize weights and apply final processing + self.post_init() + + # Model parallel + self.model_parallel = False + self.device_map = None + + @add_start_docstrings(PARALLELIZE_DOCSTRING) + def parallelize(self, device_map=None): + warnings.warn( + "`T5Model.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model" + " with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" + " `device_map` but it needs to be a dictionary module_name to device, so for instance {'encoder.block.0':" + " 0, 'encoder.block.1': 1, ...}", + FutureWarning, + ) + self.device_map = ( + get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) + if device_map is None + else device_map + ) + assert_device_map(self.device_map, len(self.encoder.block)) + self.encoder.parallelize(self.device_map) + self.decoder.parallelize(self.device_map) + self.model_parallel = True + + @add_start_docstrings(DEPARALLELIZE_DOCSTRING) + def deparallelize(self): + warnings.warn( + "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", + FutureWarning, + ) + self.encoder.deparallelize() + self.decoder.deparallelize() + self.encoder = self.encoder.to("cpu") + self.decoder = self.decoder.to("cpu") + self.model_parallel = False + self.device_map = None + torch.cuda.empty_cache() + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + self.decoder.set_input_embeddings(new_embeddings) + + def _tie_weights(self): + if self.config.tie_word_embeddings: + self._tie_or_clone_weights(self.encoder.embed_tokens, self.shared) + self._tie_or_clone_weights(self.decoder.embed_tokens, self.shared) + + def get_encoder(self): + return self.encoder + + def get_decoder(self): + return self.decoder + + def _prune_heads(self, heads_to_prune): + """ + Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base + class PreTrainedModel + """ + for layer, heads in heads_to_prune.items(): + self.encoder.layer[layer].attention.prune_heads(heads) + + @add_start_docstrings_to_model_forward(T5_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=Seq2SeqModelOutput, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.BoolTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + decoder_head_mask: Optional[torch.FloatTensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.Tensor] = None, + decoder_inputs_embeds: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.FloatTensor], Seq2SeqModelOutput]: + r""" + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, T5Model + + >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small") + >>> model = T5Model.from_pretrained("google-t5/t5-small") + + >>> input_ids = tokenizer( + ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" + ... ).input_ids # Batch size 1 + >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids # Batch size 1 + + >>> # preprocess: Prepend decoder_input_ids with start token which is pad token for T5Model. + >>> # This is not needed for torch's T5ForConditionalGeneration as it does this internally using labels arg. + >>> decoder_input_ids = model._shift_right(decoder_input_ids) + + >>> # forward pass + >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) + >>> last_hidden_states = outputs.last_hidden_state + ```""" + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask + if head_mask is not None and decoder_head_mask is None: + if self.config.num_layers == self.config.num_decoder_layers: + warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) + decoder_head_mask = head_mask + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + encoder_outputs = BaseModelOutput( + last_hidden_state=encoder_outputs[0], + hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, + attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + ) + + hidden_states = encoder_outputs[0] + + # Set device for model parallelism + if self.model_parallel: + torch.cuda.set_device(self.decoder.first_device) + hidden_states = hidden_states.to(self.decoder.first_device) + if decoder_input_ids is not None: + decoder_input_ids = decoder_input_ids.to(self.decoder.first_device) + if attention_mask is not None: + attention_mask = attention_mask.to(self.decoder.first_device) + if decoder_attention_mask is not None: + decoder_attention_mask = decoder_attention_mask.to( + self.decoder.first_device + ) + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=past_key_values, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if not return_dict: + return decoder_outputs + encoder_outputs + + return Seq2SeqModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) + + +@add_start_docstrings( + """T5 Model with a `language modeling` head on top.""", T5_START_DOCSTRING +) +class T5ForConditionalGeneration(T5PreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + "decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight", + ] + _tied_weights_keys = [ + "encoder.embed_tokens.weight", + "decoder.embed_tokens.weight", + "lm_head.weight", + ] + + def __init__(self, config: T5Config): + super().__init__(config) + self.model_dim = config.d_model + + self.shared = nn.Embedding(config.vocab_size, config.d_model) + + encoder_config = copy.deepcopy(config) + encoder_config.is_decoder = False + encoder_config.use_cache = False + encoder_config.is_encoder_decoder = False + self.encoder = T5Stack(encoder_config, self.shared) + + decoder_config = copy.deepcopy(config) + decoder_config.is_decoder = True + decoder_config.is_encoder_decoder = False + decoder_config.num_layers = config.num_decoder_layers + self.decoder = T5Stack(decoder_config, self.shared) + + self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + # Model parallel + self.model_parallel = False + self.device_map = None + + @add_start_docstrings(PARALLELIZE_DOCSTRING) + def parallelize(self, device_map=None): + warnings.warn( + "`T5ForConditionalGeneration.parallelize` is deprecated and will be removed in v5 of Transformers, you" + " should load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also" + " provide your own `device_map` but it needs to be a dictionary module_name to device, so for instance" + " {'encoder.block.0': 0, 'encoder.block.1': 1, ...}", + FutureWarning, + ) + self.device_map = ( + get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) + if device_map is None + else device_map + ) + assert_device_map(self.device_map, len(self.encoder.block)) + self.encoder.parallelize(self.device_map) + self.decoder.parallelize(self.device_map) + self.lm_head = self.lm_head.to(self.decoder.first_device) + self.model_parallel = True + + @add_start_docstrings(DEPARALLELIZE_DOCSTRING) + def deparallelize(self): + warnings.warn( + "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", + FutureWarning, + ) + self.encoder.deparallelize() + self.decoder.deparallelize() + self.encoder = self.encoder.to("cpu") + self.decoder = self.decoder.to("cpu") + self.lm_head = self.lm_head.to("cpu") + self.model_parallel = False + self.device_map = None + torch.cuda.empty_cache() + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + self.decoder.set_input_embeddings(new_embeddings) + + def _tie_weights(self): + if self.config.tie_word_embeddings: + self._tie_or_clone_weights(self.encoder.embed_tokens, self.shared) + self._tie_or_clone_weights(self.decoder.embed_tokens, self.shared) + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_output_embeddings(self): + return self.lm_head + + def get_encoder(self): + return self.encoder + + def get_decoder(self): + return self.decoder + + @add_start_docstrings_to_model_forward(T5_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.BoolTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + decoder_head_mask: Optional[torch.FloatTensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + decoder_inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ..., + config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for + labels in `[0, ..., config.vocab_size]` + + Returns: + + Examples: + + ```python + >>> from transformers import AutoTokenizer, T5ForConditionalGeneration + + >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small") + >>> model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-small") + + >>> # training + >>> input_ids = tokenizer("The walks in park", return_tensors="pt").input_ids + >>> labels = tokenizer(" cute dog the ", return_tensors="pt").input_ids + >>> outputs = model(input_ids=input_ids, labels=labels) + >>> loss = outputs.loss + >>> logits = outputs.logits + + >>> # inference + >>> input_ids = tokenizer( + ... "summarize: studies have shown that owning a dog is good for you", return_tensors="pt" + ... ).input_ids # Batch size 1 + >>> outputs = model.generate(input_ids) + >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True)) + >>> # studies have shown that owning a dog is good for you. + ```""" + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask + if head_mask is not None and decoder_head_mask is None: + if self.config.num_layers == self.config.num_decoder_layers: + warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) + decoder_head_mask = head_mask + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + # Convert encoder inputs in embeddings if needed + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + encoder_outputs = BaseModelOutput( + last_hidden_state=encoder_outputs[0], + hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, + attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + ) + + hidden_states = encoder_outputs[0] + + if self.model_parallel: + torch.cuda.set_device(self.decoder.first_device) + + if ( + labels is not None + and decoder_input_ids is None + and decoder_inputs_embeds is None + ): + # get decoder inputs from shifting lm labels to the right + decoder_input_ids = self._shift_right(labels) + + # Set device for model parallelism + if self.model_parallel: + torch.cuda.set_device(self.decoder.first_device) + hidden_states = hidden_states.to(self.decoder.first_device) + if decoder_input_ids is not None: + decoder_input_ids = decoder_input_ids.to(self.decoder.first_device) + if attention_mask is not None: + attention_mask = attention_mask.to(self.decoder.first_device) + if decoder_attention_mask is not None: + decoder_attention_mask = decoder_attention_mask.to( + self.decoder.first_device + ) + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=past_key_values, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = decoder_outputs[0] + + # Set device for model parallelism + if self.model_parallel: + torch.cuda.set_device(self.encoder.first_device) + self.lm_head = self.lm_head.to(self.encoder.first_device) + sequence_output = sequence_output.to(self.lm_head.weight.device) + + if self.config.tie_word_embeddings: + # Rescale output before projecting on vocab + # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586 + sequence_output = sequence_output * (self.model_dim**-0.5) + + lm_logits = self.lm_head(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss(ignore_index=-100) + # move labels to correct device to enable PP + labels = labels.to(lm_logits.device) + loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) + # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666 + + if not return_dict: + output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs + return ((loss,) + output) if loss is not None else output + + return Seq2SeqLMOutput( + loss=loss, + logits=lm_logits, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + head_mask=None, + decoder_head_mask=None, + decoder_attention_mask=None, + cross_attn_head_mask=None, + use_cache=None, + encoder_outputs=None, + **kwargs, + ): + # cut decoder_input_ids if past_key_values is used + if past_key_values is not None: + past_length = past_key_values[0][0].shape[2] + + # Some generation methods already pass only the last input ID + if input_ids.shape[1] > past_length: + remove_prefix_length = past_length + else: + # Default to old behavior: keep only final ID + remove_prefix_length = input_ids.shape[1] - 1 + + input_ids = input_ids[:, remove_prefix_length:] + + return { + "decoder_input_ids": input_ids, + "past_key_values": past_key_values, + "encoder_outputs": encoder_outputs, + "attention_mask": attention_mask, + "head_mask": head_mask, + "decoder_head_mask": decoder_head_mask, + "decoder_attention_mask": decoder_attention_mask, + "cross_attn_head_mask": cross_attn_head_mask, + "use_cache": use_cache, + } + + def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor): + return self._shift_right(labels) + + def _reorder_cache(self, past_key_values, beam_idx): + # if decoder past is not included in output + # speedy decoding is disabled and no need to reorder + if past_key_values is None: + logger.warning( + "You might want to consider setting `use_cache=True` to speed up decoding" + ) + return past_key_values + + reordered_decoder_past = () + for layer_past_states in past_key_values: + # get the correct batch idx from layer past batch dim + # batch dim of `past` is at 2nd position + reordered_layer_past_states = () + for layer_past_state in layer_past_states: + # need to set correct `past` for each of the four key / value states + reordered_layer_past_states = reordered_layer_past_states + ( + layer_past_state.index_select( + 0, beam_idx.to(layer_past_state.device) + ), + ) + + if reordered_layer_past_states[0].shape != layer_past_states[0].shape: + raise ValueError( + f"reordered_layer_past_states[0] shape {reordered_layer_past_states[0].shape} and layer_past_states[0] shape {layer_past_states[0].shape} mismatched" + ) + if len(reordered_layer_past_states) != len(layer_past_states): + raise ValueError( + f"length of reordered_layer_past_states {len(reordered_layer_past_states)} and length of layer_past_states {len(layer_past_states)} mismatched" + ) + + reordered_decoder_past = reordered_decoder_past + ( + reordered_layer_past_states, + ) + return reordered_decoder_past + + +@add_start_docstrings( + "The bare T5 Model transformer outputting encoder's raw hidden-states without any specific head on top.", + T5_START_DOCSTRING, +) +class T5EncoderModel(T5PreTrainedModel): + _tied_weights_keys = ["encoder.embed_tokens.weight"] + _keys_to_ignore_on_load_unexpected = [r"decoder"] + + def __init__(self, config: T5Config): + super().__init__(config) + self.shared = nn.Embedding(config.vocab_size, config.d_model) + + encoder_config = copy.deepcopy(config) + encoder_config.use_cache = False + encoder_config.is_encoder_decoder = False + self.encoder = T5Stack(encoder_config, self.shared) + + # Initialize weights and apply final processing + self.post_init() + + # Model parallel + self.model_parallel = False + self.device_map = None + + @add_start_docstrings(PARALLELIZE_DOCSTRING) + def parallelize(self, device_map=None): + warnings.warn( + "`T5EncoderModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should load" + " your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" + " `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0," + " 'block.1': 1, ...}", + FutureWarning, + ) + self.device_map = ( + get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) + if device_map is None + else device_map + ) + assert_device_map(self.device_map, len(self.encoder.block)) + self.encoder.parallelize(self.device_map) + self.model_parallel = True + + @add_start_docstrings(DEPARALLELIZE_DOCSTRING) + def deparallelize(self): + warnings.warn( + "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", + FutureWarning, + ) + self.encoder.deparallelize() + self.encoder = self.encoder.to("cpu") + self.model_parallel = False + self.device_map = None + torch.cuda.empty_cache() + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + + def _tie_weights(self): + if self.config.tie_word_embeddings: + self._tie_or_clone_weights(self.encoder.embed_tokens, self.shared) + + def get_encoder(self): + return self.encoder + + def _prune_heads(self, heads_to_prune): + """ + Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base + class PreTrainedModel + """ + for layer, heads in heads_to_prune.items(): + self.encoder.block[layer].layer[0].SelfAttention.prune_heads(heads) + + @add_start_docstrings_to_model_forward(T5_ENCODER_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=BaseModelOutput, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.FloatTensor], BaseModelOutput]: + r""" + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, T5EncoderModel + + >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small") + >>> model = T5EncoderModel.from_pretrained("google-t5/t5-small") + >>> input_ids = tokenizer( + ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" + ... ).input_ids # Batch size 1 + >>> outputs = model(input_ids=input_ids) + >>> last_hidden_states = outputs.last_hidden_state + ```""" + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + return encoder_outputs + + +@add_start_docstrings( + """ + T5 model with a sequence classification/head on top (a linear layer on top of the pooled output) e.g. for GLUE + tasks. + """, + T5_START_DOCSTRING, +) +class T5ForSequenceClassification(T5PreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + "decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight" + ] + _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] + + def __init__(self, config: T5Config): + super().__init__(config) + self.transformer = T5Model(config) + self.classification_head = T5ClassificationHead(config) + + # Initialize weights and apply final processing + self.post_init() + + self.model_parallel = False + + @add_start_docstrings_to_model_forward(T5_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=Seq2SeqSequenceClassifierOutput, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + decoder_head_mask: Optional[torch.Tensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + decoder_inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, Seq2SeqSequenceClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + Returns: + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + if labels is not None: + use_cache = False + + if input_ids is None and inputs_embeds is not None: + raise NotImplementedError( + f"Passing input embeddings is currently not supported for {self.__class__.__name__}" + ) + + # Copied from models.bart.modeling_bart.BartModel.forward different to other models, T5 automatically creates + # decoder_input_ids from input_ids if no decoder_input_ids are provided + if decoder_input_ids is None and decoder_inputs_embeds is None: + if input_ids is None: + raise ValueError( + "If no `decoder_input_ids` or `decoder_inputs_embeds` are " + "passed, `input_ids` cannot be `None`. Please pass either " + "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." + ) + decoder_input_ids = self._shift_right(input_ids) + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + head_mask=head_mask, + decoder_head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + encoder_outputs=encoder_outputs, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + + eos_mask = input_ids.eq(self.config.eos_token_id).to(sequence_output.device) + + if len(torch.unique_consecutive(eos_mask.sum(1))) > 1: + raise ValueError("All examples must have the same number of tokens.") + batch_size, _, hidden_size = sequence_output.shape + sentence_representation = sequence_output[eos_mask, :].view( + batch_size, -1, hidden_size + )[:, -1, :] + logits = self.classification_head(sentence_representation) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.config.num_labels == 1: + self.config.problem_type = "regression" + elif self.config.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.config.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct( + logits.view(-1, self.config.num_labels), labels.view(-1) + ) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return Seq2SeqSequenceClassifierOutput( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + ) + + +@add_start_docstrings( + """ + T5 Encoder Model with a token classification head on top (a linear layer on top of the hidden-states output) + e.g. for Named-Entity-Recognition (NER) tasks. + """, + T5_START_DOCSTRING, +) +class T5ForTokenClassification(T5PreTrainedModel): + _tied_weights_keys = ["transformer.encoder.embed_tokens.weight"] + + def __init__(self, config: T5Config): + super().__init__(config) + self.num_labels = config.num_labels + + self.transformer = T5EncoderModel(config) + self.dropout = nn.Dropout(config.classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(T5_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=TokenClassifierOutput, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + Returns: + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + hidden_states = self.dropout(hidden_states) + logits = self.classifier(hidden_states) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits, outputs[2:-1]) + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + T5 Model with a span classification head on top for extractive question-answering tasks like SQuAD (linear layers + on top of the hidden-states output to compute `span start logits` and `span end logits`). + """, + T5_START_DOCSTRING, +) +class T5ForQuestionAnswering(T5PreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + "decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight" + ] + _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] + + def __init__(self, config: T5Config): + super().__init__(config) + self.model_dim = config.d_model + + self.shared = nn.Embedding(config.vocab_size, config.d_model) + + encoder_config = copy.deepcopy(config) + encoder_config.is_decoder = False + encoder_config.use_cache = False + encoder_config.is_encoder_decoder = False + self.encoder = T5Stack(encoder_config, self.shared) + + decoder_config = copy.deepcopy(config) + decoder_config.is_decoder = True + decoder_config.is_encoder_decoder = False + decoder_config.num_layers = config.num_decoder_layers + self.decoder = T5Stack(decoder_config, self.shared) + + self.num_labels = config.num_labels + self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + self.model_parallel = False + + def get_input_embeddings(self): + return self.shared + + def set_input_embeddings(self, new_embeddings): + self.shared = new_embeddings + self.encoder.set_input_embeddings(new_embeddings) + self.decoder.set_input_embeddings(new_embeddings) + + def _tie_weights(self): + if self.config.tie_word_embeddings: + self._tie_or_clone_weights(self.encoder.embed_tokens, self.shared) + self._tie_or_clone_weights(self.decoder.embed_tokens, self.shared) + + def get_encoder(self): + return self.encoder + + def get_decoder(self): + return self.decoder + + @add_start_docstrings_to_model_forward(T5_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=Seq2SeqQuestionAnsweringModelOutput, config_class=_CONFIG_FOR_DOC + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.BoolTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + decoder_head_mask: Optional[torch.FloatTensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, + start_positions: Optional[torch.LongTensor] = None, + end_positions: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + decoder_inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.FloatTensor], Seq2SeqQuestionAnsweringModelOutput]: + r""" + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence + are not taken into account for computing the loss. + Returns: + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + if start_positions is not None and end_positions is not None: + use_cache = False + + # Copied from models.bart.modeling_bart.BartModel.forward + # different to other models, T5 automatically creates decoder_input_ids from + # input_ids if no decoder_input_ids are provided + if decoder_input_ids is None and decoder_inputs_embeds is None: + if input_ids is None: + raise ValueError( + "If no `decoder_input_ids` or `decoder_inputs_embeds` are " + "passed, `input_ids` cannot be `None`. Please pass either " + "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." + ) + decoder_input_ids = self._shift_right(input_ids) + + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask + if head_mask is not None and decoder_head_mask is None: + if self.config.num_layers == self.config.num_decoder_layers: + warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) + decoder_head_mask = head_mask + + # Encode if needed (training, first prediction pass) + if encoder_outputs is None: + encoder_outputs = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + encoder_outputs = BaseModelOutput( + last_hidden_state=encoder_outputs[0], + hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, + attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + ) + + hidden_states = encoder_outputs[0] + + # Decode + decoder_outputs = self.decoder( + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + inputs_embeds=decoder_inputs_embeds, + past_key_values=None, + encoder_hidden_states=hidden_states, + encoder_attention_mask=attention_mask, + head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = decoder_outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1).to(start_logits.device) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1).to(end_logits.device) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + decoder_outputs[1:] + encoder_outputs + return ((total_loss,) + output) if total_loss is not None else output + + return Seq2SeqQuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, + encoder_hidden_states=encoder_outputs.hidden_states, + encoder_attentions=encoder_outputs.attentions, + ) diff --git a/ixformer_sdk/contrib/vllm/__init__.py b/ixformer_sdk/contrib/vllm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/contrib/vllm/layers/__init__.py b/ixformer_sdk/contrib/vllm/layers/__init__.py new file mode 100644 index 00000000..79eb0e96 --- /dev/null +++ b/ixformer_sdk/contrib/vllm/layers/__init__.py @@ -0,0 +1,30 @@ +from .llama import forward_smoothquant +from .mixtral import mixtral_decoder_layer_forward + +SUPPORT_REPLACE_METHOD = { + "llama": forward_smoothquant, +} + +SUPPORT_REPLACE_LAYER = { + "llama": None, +} + + +def get_replace_forward(name: str): + try: + method = SUPPORT_REPLACE_METHOD[name] + except: + raise ValueError( + f"Only support replace names: {SUPPORT_REPLACE_METHOD.keys()}, but got {name}" + ) + return method + + +def get_replace_layer(name: str): + try: + layer = SUPPORT_REPLACE_LAYER[name] + except: + raise ValueError( + f"Only support replace names: {SUPPORT_REPLACE_LAYER.keys()}, but got {name}" + ) + return layer diff --git a/ixformer_sdk/contrib/vllm/layers/llama.py b/ixformer_sdk/contrib/vllm/layers/llama.py new file mode 100644 index 00000000..0e670080 --- /dev/null +++ b/ixformer_sdk/contrib/vllm/layers/llama.py @@ -0,0 +1,100 @@ +from typing import Any, Dict, Iterable, List, Optional, Tuple + +import torch +from transformers import LlamaConfig + +from vllm.attention import AttentionMetadata +from vllm.config import CacheConfig +from vllm.distributed import tensor_model_parallel_all_reduce +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.models.llama import LlamaDecoderLayer as VllmLlamaDecoderLayer + +import vllm._custom_ops as ops +# from ..overlap_comm import DecoderLayerOverlapComm, get_overlap_linear_method + + +# This method is needed for support smoothquant no overlap forward +def forward_smoothquant( + input_ids: Optional[torch.Tensor], + positions: torch.Tensor, + kv_caches: List[torch.Tensor], + attn_metadata: AttentionMetadata, + inputs_embeds: Optional[torch.Tensor] = None, + self = None, # will be set by partial +) -> torch.Tensor: + dtype = self.dtype + + def forward_smoothquant_mlp(self,x,scales): + # gate_up_proj + # Int8 Matrix multiply. + bias = self.gate_up_proj.bias if not self.gate_up_proj.skip_bias_add else None + gate_up = ops.w8a8(x, self.gate_up_proj.weight, scales, self.gate_up_proj.weight_scales, dtype) + if bias: + gate_up += bias + + # act_fun + x, scales = ops.silu_and_mul_smoothquant(gate_up, self.down_proj.smooth_scales) + + # down_proj + output_parallel = ops.w8a8(x, self.down_proj.weight, scales, self.down_proj.weight_scales, dtype) + if self.down_proj.reduce_results and self.down_proj.tp_size > 1: + output = tensor_model_parallel_all_reduce(output_parallel) + else: + output = output_parallel + + if not self.down_proj.skip_bias_add: + output = output + self.down_proj.bias if self.down_proj.bias is not None else output + + return output + + def forward_smoothquant_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: AttentionMetadata, + scales: torch.Tensor, + ) -> torch.Tensor: + # qkv proj + bias = self.qkv_proj.bias if not self.qkv_proj.skip_bias_add else None + + qkv = ops.w8a8(hidden_states, self.qkv_proj.weight, scales, self.qkv_proj.weight_scales, dtype) + if bias: + qkv += bias + + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, kv_cache, attn_metadata) + output, _ = self.o_proj(attn_output) # TODO + return output + + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.get_input_embeddings(input_ids) + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + if residual is None: + residual = hidden_states + hidden_states, scales = ops.rms_norm_smoothquant(hidden_states,layer.input_layernorm.weight,layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales) + else: + hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.input_layernorm.weight, layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales) + + hidden_states = forward_smoothquant_attn( + layer.self_attn, + positions=positions, + hidden_states=hidden_states, + kv_cache=kv_caches[i], + attn_metadata=attn_metadata, + scales=scales, + ) + + # Fully Connected + hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.post_attention_layernorm.weight, layer.post_attention_layernorm.variance_epsilon, layer.mlp.gate_up_proj.smooth_scales) + + hidden_states = forward_smoothquant_mlp(layer.mlp, hidden_states, scales) + + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + diff --git a/ixformer_sdk/contrib/vllm/layers/mixtral.py b/ixformer_sdk/contrib/vllm/layers/mixtral.py new file mode 100644 index 00000000..be4dd003 --- /dev/null +++ b/ixformer_sdk/contrib/vllm/layers/mixtral.py @@ -0,0 +1,331 @@ +import functools +from typing import Dict, Optional, Tuple + +import ixformer.inference.functions as ixf +import torch + + +def mixtral_decoder_layer_forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata, + residual: Optional[torch.Tensor], +) -> torch.Tensor: + if self.use_int_w8a8: + return w8a8_forward( + self, positions, hidden_states, kv_cache, attn_metadata, residual + ) + else: + return original_forward( + self, positions, hidden_states, kv_cache, attn_metadata, residual + ) + + +def original_forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata, + residual: Optional[torch.Tensor], +) -> torch.Tensor: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + kv_cache=kv_cache, + attn_metadata=attn_metadata, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.block_sparse_moe(hidden_states) + return hidden_states, residual + + +def dynamic_scaled_int8_quant(x): + m, k = x.shape + i8_x = x.new_empty([m, k], dtype=torch.int8, device="cuda") + i8_scales = torch.empty([m], dtype=torch.float32, device="cuda") + ixf.dynamic_scaled_int8_quant(i8_x, x, i8_scales) + return i8_x, i8_scales + + +def dynamic_w8a8(x, i8_weight, weight_scale): + i8_x, i8_scale = dynamic_scaled_int8_quant(x) + m, k = x.shape + k, n = i8_weight.shape + output = x.new_empty([m, n], dtype=x.dtype, device="cuda") + ixf.w8a8( + i8_x, + i8_weight.transpose(0, 1), + i8_scale, + weight_scale, + output=output, + out_dtype=x.dtype, + ) + return output + + +def fused_rms_norm_quant_linear( + self, + hidden_states, + ln_weight, + eps, + linear_weight, + linear_weight_scale, + residual=None, +): + # lower rouge + # if residual is None: + # residual = hidden_states + # i8_hidden_states, _, i8_scales = ixf.residual_rms_norm_dynamic_int8( + # input=hidden_states, + # weight=ln_weight, + # residual=None, + # eps=eps, + # ) + # else: + # i8_hidden_states, residual, i8_scales = ixf.residual_rms_norm_dynamic_int8( + # input=hidden_states, + # weight=ln_weight, + # residual=residual, + # eps=eps, + # ) + + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + i8_hidden_states, i8_scales = dynamic_scaled_int8_quant(hidden_states) + + qkv = hidden_states.new_empty(hidden_states.shape[0], linear_weight.shape[1]) + ixf.w8a8( + i8_hidden_states, + linear_weight.transpose(0, 1), + i8_scales, + linear_weight_scale, + output=qkv, + out_dtype=hidden_states.dtype, + ) + return qkv, residual + + +def attention(qkv, positions, kv_cache, attn_metadata, self_attn): + q, k, v = qkv.split( + [self_attn.q_size, self_attn.kv_size, self_attn.kv_size], dim=-1 + ) + q, k = self_attn.rotary_emb(positions, q, k) + attn_output = self_attn.attn(q, k, v, kv_cache, attn_metadata) + return attn_output + + +def fused_rms_norm_attention( + self, + hidden_states, + ln_weight, + eps, + positions, + kv_cache, + attn_metadata, + self_attn, + residual=None, +): + hidden_states, residual = fused_rms_norm_quant_linear( + self, + hidden_states, + ln_weight, + eps, + self_attn.qkv_proj.weight, + self_attn.qkv_proj.weight_scale, + residual, + ) + hidden_states = attention( + hidden_states, positions, kv_cache, attn_metadata, self_attn + ) + + hidden_states = dynamic_w8a8( + hidden_states, self_attn.o_proj.weight, self_attn.o_proj.weight_scale + ) + # hidden_states,_ = self_attn.o_proj(hidden_states) # quant+linear+allreduce + return hidden_states, residual + + +def w8a8_forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata, + residual: Optional[torch.Tensor], +) -> torch.Tensor: + + # qkv,_ = self.self_attn.qkv_proj(hidden_states) + hidden_states, residual = fused_rms_norm_attention( + self, + hidden_states, + self.input_layernorm.weight, + self.input_layernorm.variance_epsilon, + positions, + kv_cache, + attn_metadata, + self.self_attn, + residual, + ) + + # allreduce + tp_size = self.block_sparse_moe.experts.tp_size + if tp_size > 1: + from vllm.distributed import tensor_model_parallel_all_reduce + + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + + # rms norm + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + # moe + hidden_states = fused_moe( + hidden_states, + self.block_sparse_moe.gate.weight, + top_k=self.block_sparse_moe.experts.top_k, + w1=self.block_sparse_moe.experts.w13_weight, + w2=self.block_sparse_moe.experts.w2_weight, + w1_scale=self.block_sparse_moe.experts.w13_weight_scale, + w2_scale=self.block_sparse_moe.experts.w2_weight_scale, + ) + + # allreduce + if tp_size > 1: + from vllm.distributed import tensor_model_parallel_all_reduce + + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + + return hidden_states, residual + + +def fused_experts(hidden_states, router_logits, top_k, w1, w2, w1_scale, w2_scale): + + """ + Args: + hidden_states: (num_tokens, k) dtype + router_logits: (num_tokens, num_experts) torch.float32 + top_k int + w1: (num_experts, 2n, k) torch.int8 + w2: (num_experts, k, n) torch.int8 + w1_scale: (num_experts, 2n) torch.float32 + w2_scale: (num_experts, k) torch.float32 + Returns + final_hidden_states: (num_tokens, k) dtype + """ + + # topk_weight: (num_tokens, top_k) torch.float32 + # topk_ids: (num_tokens, top_k) torch.int32 + topk_weight, topk_ids = ixf.moe_topk_softmax( + gating_output=router_logits, + topk=top_k, + renormalize=True, + ) + + dtype = hidden_states.dtype + num_tokens, num_experts = router_logits.shape + expand_tokens = num_tokens * top_k + + ( + src_to_dst, + sorted_token_ids, + expert_sizes_gpu, + expert_sizes_cpu, + ) = ixf.moe_compute_token_index( + topk_ids=topk_ids, + num_experts=num_experts, + ) + expert_sizes_cpu = expert_sizes_gpu.cpu() + + # expand + reorder + quant + # i8_hidden_states: (expand_tokens, k) torch.int8 + i8_hidden_states, a_scale = ixf.moe_expand_input_dynamic_scaled_int8( + hidden_states=hidden_states, + dst_to_src=sorted_token_ids, + dst_tokens=expand_tokens, + topk=top_k, + src_to_dst=src_to_dst, + topk_ids=None, # use smooth quant + smooth_scales=None, # use smooth quant + ) + + # w8a8 group gemm 1 + # pt_output_1: (expand_tokens, 2n) dtype + pt_output_1 = ixf.moe_w8a8_group_gemm( + input=i8_hidden_states, + weight=w1, + i_scales=a_scale, + w_scales=w1_scale, + output_dtype=dtype, + tokens_per_experts=expert_sizes_cpu, + dst_to_src=None, + format="TN", + ) + + # act + quant + # pt_output_2: (expand_tokens, n) torch.int8 + pt_output_2, a2_scale = ixf.activation_dynamic_scaled_int8( + input=pt_output_1, + bias=None, # add gemm bias + smooth_scales=None, # use smooth quant + dst_to_src=sorted_token_ids, + topk_ids=None, # add gemm bias or use smooth quant + act_type="swiglu", + ) + + # w8a8 group gemm 2 + reorder + # pt_output_3: (expand_tokens, k) dtype + pt_output_3 = ixf.moe_w8a8_group_gemm( + input=pt_output_2, + weight=w2, + i_scales=a2_scale, + w_scales=w2_scale, + output_dtype=dtype, + tokens_per_experts=expert_sizes_cpu, + dst_to_src=sorted_token_ids, + format="TN", + ) + + # mul + reduce_sum + # final_hidden_states: (num_tokens, k) + final_hidden_states = ixf.moe_output_reduce_sum( + input=pt_output_3.view(num_tokens, top_k, -1), + topk_weight=topk_weight, + ) + + return final_hidden_states + + +def fused_moe(hidden_states, gate_weight, top_k, w1, w2, w1_scale, w2_scale): + orig_shape = hidden_states.shape + hidden_size = hidden_states.shape[-1] + + hidden_states = hidden_states.view(-1, hidden_size) + + # router_logits: (num_tokens, n_experts) + # gate_weight: fp16 + router_logits = ixf.linear(hidden_states, gate_weight) + router_logits = router_logits.to(torch.float32) + + final_hidden_states = fused_experts( + hidden_states, + router_logits, + top_k, + w1, + w2, + w1_scale, + w2_scale, + ) + + return final_hidden_states.view(orig_shape) diff --git a/ixformer_sdk/contrib/vllm/quantize/__init__.py b/ixformer_sdk/contrib/vllm/quantize/__init__.py new file mode 100644 index 00000000..f2b68c4a --- /dev/null +++ b/ixformer_sdk/contrib/vllm/quantize/__init__.py @@ -0,0 +1,14 @@ +from .smoothquant import smoothquant_prepare_quantize,smoothquant_export_quantized_weights +from .w8a16 import w8a16_prepare_quantize,w8a16_export_quantized_weights + +SUPPORT_METHOD = { + "smoothquant": [smoothquant_prepare_quantize,smoothquant_export_quantized_weights], + "w8a16": [w8a16_prepare_quantize,w8a16_export_quantized_weights], +} + +def get_quantize_method(method_name:str): + try: + method = SUPPORT_METHOD[method_name] + except: + raise ValueError(f"Only support quantization methods: {SUPPORT_METHOD.keys()}, but got {method_name}") + return method \ No newline at end of file diff --git a/ixformer_sdk/contrib/vllm/quantize/smoothquant.py b/ixformer_sdk/contrib/vllm/quantize/smoothquant.py new file mode 100644 index 00000000..63960e25 --- /dev/null +++ b/ixformer_sdk/contrib/vllm/quantize/smoothquant.py @@ -0,0 +1,407 @@ +import os + +import torch + + +def smoothquant_prepare_quantize(self, quant_params={}): + model = self.model_runner.model + + def update_act_scales(act_scales, x): + # 动态统计每次输入的最大值 + hidden_dim = x.shape[-1] + x = x.view(-1, hidden_dim).abs().detach() + # [k] + comming_max = torch.max(x, dim=0, keepdim=True)[0].float() + + if act_scales is None: + act_scales = comming_max + else: + act_scales = torch.max(act_scales, comming_max) + return act_scales + + from functools import partial + + from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, + ) + + def new_forward(input_, m, raw_forward): + if not hasattr(m, "act_scales"): + m.act_scales = None + m.act_scales = update_act_scales(m.act_scales, input_) + return raw_forward(input_) + + for name, m in model.named_modules(): + if ( + isinstance(m, QKVParallelLinear) + or isinstance(m, RowParallelLinear) + or isinstance(m, MergedColumnParallelLinear) + or isinstance(m, ColumnParallelLinear) + ): + m.forward = partial(new_forward, m=m, raw_forward=m.forward) + +def smoothquant_export_quantized_weights(self, save_path, quant_params={}): + gb_per_file = quant_params.get("filesize_limit", None) + smooth_alpha = quant_params.get("smooth_alpha", 0.5) + dynamic_quant_type = quant_params.get("dynamic_quant_type", "gpu") + assert dynamic_quant_type in ["gpu","cpu","kernel"] + if self.rank == 0: + print(f"set smooth_alpha={smooth_alpha}") + print(f"use quantize weight type: {dynamic_quant_type}") + + import ixformer._C as ops + def per_token_quant_8bit(weight): + # weight: [m,k] + dtype = weight.dtype + i8_weight = weight + scale = i8_weight.abs().max(dim=-1, keepdim=True)[0] / 127 + i8_weight = i8_weight / scale.to(dtype) + i8_weight = torch.clamp(torch.round(i8_weight), -128, 127).to(torch.int8) + return i8_weight, scale.float() + + def smooth_quant_weight_gpu_cpu(weight, act_scale, alpha=0.5, device="cpu"): + device = torch.device("cpu") if device == "cpu" else weight.device + ori_dtype = weight.dtype + # [1, k] + act_scale = act_scale.float().to(device).view(1, -1) + weight = weight.to(device) + # [1, k] + weight_scale = weight.abs().max(dim=0, keepdim=True)[0].float() + if alpha == -1: + smooth_scales = torch.ones_like(act_scale) + else: + smooth_scales = act_scale.pow(alpha) / weight_scale.pow(1 - alpha).clamp( + min=1e-5 + ) + weight = weight * smooth_scales.to(ori_dtype) + i8_weight, weight_scales = per_token_quant_8bit(weight) + # 为了可以使用 input * smooth_scales + if alpha == -1: + smooth_scales = torch.ones_like(act_scale) + else: + smooth_scales = weight_scale.pow(1 - alpha) / act_scale.pow(alpha).clamp( + min=1e-5 + ) + return i8_weight, weight_scales, smooth_scales.to(ori_dtype) + + def smooth_quant_weight_kernel(weight, act_scale, alpha=0.5): + output = torch.zeros_like(weight,dtype=torch.int8) + weight_scales = torch.zeros(weight.shape[:-1],dtype=torch.float, device=weight.device) + weight_max = torch.zeros(weight.shape[-1],dtype=torch.float, device=weight.device) + smooth_scales = torch.zeros(weight.shape[-1],dtype=weight.dtype, device=weight.device) + ops.infer.weight_quant_smoothquant( + weight, act_scale, alpha, output, weight_scales, smooth_scales, weight_max + ) + return output, weight_scales.view(-1,1), smooth_scales.view(1,-1) + + def smooth_quant_weight(weight, act_scale, alpha=0.5): + if dynamic_quant_type == "kernel": + return smooth_quant_weight_kernel(weight,act_scale,alpha) + else: + return smooth_quant_weight_gpu_cpu(weight,act_scale,alpha,dynamic_quant_type) + + model = self.model_runner.model + + from vllm.distributed import ( + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, + get_tensor_model_parallel_world_size + ) + from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, + ) + from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, + ) + from vllm.model_executor.models.falcon import FalconForCausalLM + + for name, m in model.named_modules(): + if isinstance(m, VocabParallelEmbedding): + # weight shape: [vocab_size // tp, embedding_dim] + weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0) + weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous() + if self.is_driver_worker: + m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False) + print(f"merged: {name}, shape={m.weight.shape}") + + elif isinstance(m, ParallelLMHead): + # weight shape: [vocab_size // tp, embedding_dim] + # bias shape: [vocab_size // tp] + if m.bias is not None: + bias = tensor_model_parallel_all_gather(m.bias, dim=0) + bias = bias[:m.org_vocab_size].contiguous() + if self.is_driver_worker: + m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False) + + weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0) + weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous() + if self.is_driver_worker: + m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False) + print(f"merged: {name}, shape={m.weight.shape}") + + elif isinstance(m, QKVParallelLinear): + # weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size] + # bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp] + if self.parallel_config.world_size > 1: + total_q_hidden_size = m.total_num_heads * m.head_size + partial_q_hidden_size = m.num_heads * m.head_size + total_kv_hidden_size = m.total_num_kv_heads * m.head_size + partial_kv_hidden_size = m.num_kv_heads * m.head_size + + if m.bias is not None: + # TODO do not support padding.. + bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size) + q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size] + k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\ + [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size] + v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\ + [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size] + + q_bias[:] = m.bias[:partial_q_hidden_size] + k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size] + v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:] + + bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr) + if self.is_driver_worker: + m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False) + + q_tensor = m.weight.new_zeros(m.total_num_heads * m.head_size, m.weight.shape[1]) + k_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1]) + v_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1]) + + q_in_weight = m.weight[:-m.num_kv_heads * m.head_size * 2] + k_in_weight = m.weight[-m.num_kv_heads * m.head_size * 2:-m.num_kv_heads * m.head_size] + v_in_weight = m.weight[-m.num_kv_heads * m.head_size:] + + if getattr(m,"start_idx",None) is not None: + start_idx = getattr(m,"start_idx") + weight_end_idx = m.num_heads * m.head_size if not getattr(m,"is_padding") else (m.num_heads - 1) * m.head_size + end_idx = start_idx + weight_end_idx + else: + start_idx = self.rank * m.num_heads * m.head_size + weight_end_idx = m.num_heads * m.head_size + end_idx = start_idx + weight_end_idx + assert q_tensor[start_idx:end_idx,:].shape == q_in_weight[:weight_end_idx, :].shape + q_tensor[start_idx:end_idx,:] = q_in_weight[:weight_end_idx, :] + + if m.num_kv_head_replicas > 1: + if self.rank % m.num_kv_head_replicas == 0: + rank = self.rank // m.num_kv_head_replicas + k_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = k_in_weight + v_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = v_in_weight + else: + k_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = k_in_weight + v_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = v_in_weight + + q_tensor = tensor_model_parallel_all_reduce(q_tensor) + k_tensor = tensor_model_parallel_all_reduce(k_tensor) + v_tensor = tensor_model_parallel_all_reduce(v_tensor) + + if isinstance(model, FalconForCausalLM): + num_query_heads_per_kv_head = ( + m.total_num_heads // m.total_num_kv_heads + ) + q_tensor = q_tensor.view( + m.total_num_kv_heads, + num_query_heads_per_kv_head, + m.head_size, + -1, + ) + k_tensor = k_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1) + v_tensor = v_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1) + weight_tensor = torch.cat( + [q_tensor, k_tensor, v_tensor], dim=1 + ).view(-1, m.hidden_size) + else: + weight_tensor = torch.cat([q_tensor, k_tensor, v_tensor]) + assert ( + weight_tensor.shape[0] + == total_q_hidden_size + total_kv_hidden_size * 2 + ) + assert weight_tensor.shape[1] == m.hidden_size + + else: + weight_tensor = m.weight + if m.bias is not None and self.is_driver_worker: + m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False) + + if self.is_driver_worker: + i8_weight, weight_scales, smooth_scales = smooth_quant_weight( + weight_tensor, m.act_scales, smooth_alpha + ) + m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False) + m.weight_scales = torch.nn.Parameter( + weight_scales.cpu(), requires_grad=False + ) + m.smooth_scales = torch.nn.Parameter( + smooth_scales.cpu(), requires_grad=False + ) + print(f"Quantized: {name}") + + elif isinstance(m, MergedColumnParallelLinear): + if self.parallel_config.world_size > 1: + # weight shape: [intermediate_size // tp * 2, hidden_size] + # bias shape: [intermediate_size // tp * 2] + output_sizes = m.output_sizes + output_size = sum(output_sizes) + partial_output_sizes = [ + i // self.parallel_config.world_size for i in output_sizes + ] + + if m.bias is not None: + index_start = 0 + partial_index_start = 0 + bias_tenosr = m.bias.new_zeros(output_size) + for i in range(len(output_sizes)): + index_out = index_start + output_sizes[i] + sub_bias_tensor = bias_tenosr[index_start:index_out] + partial_size = partial_output_sizes[i] + sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size] + + index_start += output_sizes[i] + partial_index_start += partial_size + bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr) + if self.is_driver_worker: + m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False) + + weight_tensor = m.weight.new_zeros(output_size, m.input_size) + + idx_out_start = 0 + idx_partial_satrt = 0 + for i in range(len(output_sizes)): + idx_out_end = idx_out_start + output_sizes[i] + sub_weight_tensor = weight_tensor[idx_out_start:idx_out_end] + partial_size = partial_output_sizes[i] + sub_weight_tensor[ + self.rank * partial_size : (self.rank + 1) * partial_size + ] = m.weight[idx_partial_satrt : idx_partial_satrt + partial_size] + + idx_out_start += output_sizes[i] + idx_partial_satrt += partial_size + weight_tensor = tensor_model_parallel_all_reduce(weight_tensor) + else: + weight_tensor = m.weight + if m.bias is not None and self.is_driver_worker: + m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False) + + if self.is_driver_worker: + i8_weight, weight_scales, smooth_scales = smooth_quant_weight( + weight_tensor, m.act_scales, smooth_alpha + ) + m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False) + m.weight_scales = torch.nn.Parameter( + weight_scales.cpu(), requires_grad=False + ) + m.smooth_scales = torch.nn.Parameter( + smooth_scales.cpu(), requires_grad=False + ) + print(f"Quantized: {name}") + + elif isinstance(m, ColumnParallelLinear): + # weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4 + # bias shape: [some_dim // tp] + if m.bias is not None: + bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0) + if self.is_driver_worker: + m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False) + + weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0) + + if self.is_driver_worker: + i8_weight, weight_scales, smooth_scales = smooth_quant_weight( + weight_tensor, m.act_scales, smooth_alpha + ) + m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False) + m.weight_scales = torch.nn.Parameter( + weight_scales.cpu(), requires_grad=False + ) + m.smooth_scales = torch.nn.Parameter( + smooth_scales.cpu(), requires_grad=False + ) + print(f"Quantized: {name}") + + elif isinstance(m, RowParallelLinear): + # weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size + # bias shape: [hidden_size] + if m.bias is not None: + bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1) + if self.is_driver_worker: + m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False) + + if getattr(m,"start_idx", None) is not None: + start_idx = getattr(m,"start_idx") + end_idx = start_idx + (m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size)) + weight_end_idx = m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size) + else: + start_idx = m.input_size_per_partition * self.rank + end_idx = start_idx + m.input_size_per_partition + weight_end_idx = m.input_size_per_partition + + act_scales = m.act_scales.new_zeros(m.input_size) + assert act_scales[start_idx:end_idx].shape == m.act_scales.view(-1)[:weight_end_idx].shape + act_scales[start_idx:end_idx] = m.act_scales.view(-1)[:weight_end_idx] + act_scales = tensor_model_parallel_all_reduce(act_scales) + m.act_scales = act_scales + + weight_tensor = m.weight.new_zeros(m.weight.shape[0],m.input_size) + assert weight_tensor[:,start_idx:end_idx].shape == m.weight[:,:weight_end_idx].shape + weight_tensor[:,start_idx:end_idx] = m.weight[:,:weight_end_idx] + weight_tensor = tensor_model_parallel_all_reduce(weight_tensor) + + if self.is_driver_worker: + i8_weight, weight_scales, smooth_scales = smooth_quant_weight( + weight_tensor, m.act_scales, smooth_alpha + ) + smooth_scales = smooth_scales.view(1,-1) + m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False) + m.weight_scales = torch.nn.Parameter( + weight_scales.cpu(), requires_grad=False + ) + m.smooth_scales = torch.nn.Parameter( + smooth_scales.cpu(), requires_grad=False + ) + print(f"Quantized: {name}") + else: + pass + + torch.cuda.empty_cache() + + # save weights + if self.is_driver_worker: + from safetensors.torch import save_file + + tensors = {} + saved = False + count = 0 + size_in_bytes = 0 + + tensors = {} + for name, weight in model.named_parameters(): + if "act_scales" in name: + continue + # skip lm_head_weight if needed.. + if "lm_head" in name and model.config.tie_word_embeddings: + continue + tensors[name] = weight + + saved = False + if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024: + weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6))) + save_file(tensors, weight_path) + print(f"The quantified weights were successfully saved in {weight_path}.") + tensors.clear() + saved = True + count += 1 + size_in_bytes = 0 + + if not saved: + weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6))) + save_file(tensors, weight_path) + print(f"The quantified weights were successfully saved in {weight_path}.") \ No newline at end of file diff --git a/ixformer_sdk/contrib/vllm/quantize/w8a16.py b/ixformer_sdk/contrib/vllm/quantize/w8a16.py new file mode 100644 index 00000000..8bbeeec3 --- /dev/null +++ b/ixformer_sdk/contrib/vllm/quantize/w8a16.py @@ -0,0 +1,233 @@ +import os + +import torch + +def w8a16_prepare_quantize(self, quant_params={}): + # We need do nothing in here + pass + + +def w8a16_export_quantized_weights(self, save_path, quant_params={}): + gb_per_file = quant_params.get("filesize_limit", None) + int8_min = -127 + + def w8a16_quantization(weight): + # all weights should be [output,input], otherwise, we may get an wrong weight and scale... + scale = torch.abs(weight).max(dim=-1)[0] / 127.0 + int8_weight = torch.clamp(weight / scale.view(-1,1),min=int8_min,max=127).to(torch.int8).contiguous() + scale = scale.view(1,-1).contiguous() + return int8_weight, scale + + + model = self.model_runner.model + + from vllm.distributed.communication_op import ( + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, + ) + from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, + ) + from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, + ) + + for name, m in model.named_modules(): + if isinstance(m, VocabParallelEmbedding): + # weight shape: [vocab_size // tp, embedding_dim] + weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0) + weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous() + if self.is_driver_worker: + m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False) + print(f"merged: {name}, shape={m.weight.shape}") + + elif isinstance(m, ParallelLMHead): + # weight shape: [vocab_size // tp, embedding_dim] + # bias shape: [vocab_size // tp] + if m.bias is not None: + bias = tensor_model_parallel_all_gather(m.bias, dim=0) + bias = bias[:m.org_vocab_size].contiguous() + if self.is_driver_worker: + m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False) + + weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0) + weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous() + if self.is_driver_worker: + m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False) + print(f"merged: {name}, shape={m.weight.shape}") + + elif isinstance(m, QKVParallelLinear): + # weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size] + # bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp] + if self.parallel_config.world_size > 1: + total_q_hidden_size = m.total_num_heads * m.head_size + partial_q_hidden_size = m.num_heads * m.head_size + total_kv_hidden_size = m.total_num_kv_heads * m.head_size + partial_kv_hidden_size = m.num_kv_heads * m.head_size + + if m.bias is not None: + bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size) + q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size] + k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\ + [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size] + v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\ + [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size] + + q_bias[:] = m.bias[:partial_q_hidden_size] + k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size] + v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:] + + bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr) + if self.is_driver_worker: + m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False) + + weight_tensor = m.weight.new_zeros( + total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size + ) + + q_tensor = weight_tensor[:total_q_hidden_size, :] + q_tensor = q_tensor[self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size] + + k_tensor = weight_tensor[total_q_hidden_size : total_q_hidden_size + total_kv_hidden_size] + k_tensor = k_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size] + + v_tensor = weight_tensor[total_q_hidden_size + total_kv_hidden_size :] + v_tensor = v_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size] + + q_tensor[:, :] = m.weight[: partial_q_hidden_size, :] + k_tensor[:, :] = m.weight[partial_q_hidden_size : partial_q_hidden_size + partial_kv_hidden_size, :] + v_tensor[:, :] = m.weight[partial_q_hidden_size + partial_kv_hidden_size : , :] + + weight_tensor = tensor_model_parallel_all_reduce(weight_tensor) + else: + weight_tensor = m.weight + if m.bias is not None and self.is_driver_worker: + m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False) + + int8_weight, weight_scales = w8a16_quantization(weight_tensor) + + if self.is_driver_worker: + m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False) + m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False) + print(f"Quantized: {name}") + + elif isinstance(m, MergedColumnParallelLinear): + # weight shape: [intermediate_size // tp * 2, hidden_size] + # bias shape: [intermediate_size // tp * 2] + if self.parallel_config.world_size > 1: + output_sizes = m.output_sizes + output_size = sum(output_sizes) + partial_output_sizes = [ + i // self.parallel_config.world_size for i in output_sizes + ] + + if m.bias is not None: + index_start = 0 + partial_index_start = 0 + bias_tenosr = m.bias.new_zeros(output_size) + for i in range(len(output_sizes)): + index_out = index_start + output_sizes[i] + sub_bias_tensor = bias_tenosr[index_start:index_out] + partial_size = partial_output_sizes[i] + sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size] + + index_start += output_sizes[i] + partial_index_start += partial_size + bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr) + if self.is_driver_worker: + m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False) + + weight_tensor = m.weight.new_zeros(output_size, m.input_size) + + index_start = 0 + partial_index_start = 0 + for i in range(len(output_sizes)): + index_out = index_start + output_sizes[i] + sub_weight_tensor = weight_tensor[index_start:index_out] + partial_size = partial_output_sizes[i] + sub_weight_tensor[self.rank * partial_size : (self.rank + 1) * partial_size] = m.weight[partial_index_start : partial_index_start + partial_size] + + index_start += output_sizes[i] + partial_index_start += partial_size + weight_tensor = tensor_model_parallel_all_reduce(weight_tensor) + else: + weight_tensor = m.weight + if m.bias is not None and self.is_driver_worker: + m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False) + + int8_weight, weight_scales = w8a16_quantization(weight_tensor) + + if self.is_driver_worker: + m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False) + m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False) + print(f"Quantized: {name}") + + elif isinstance(m, ColumnParallelLinear): + # weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4 + # bias shape: [some_dim // tp] + if m.bias is not None: + bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0) + if self.is_driver_worker: + m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False) + + weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0) + + int8_weight, weight_scales = w8a16_quantization(weight_tensor) + + if self.is_driver_worker: + m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False) + m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False) + print(f"Quantized: {name}") + + elif isinstance(m, RowParallelLinear): + # weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size + # bias shape: [hidden_size] + if m.bias is not None: + bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1) + if self.is_driver_worker: + m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False) + + weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=-1) + int8_weight, weight_scales = w8a16_quantization(weight_tensor) + + if self.is_driver_worker: + m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False) + m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False) + print(f"Quantized: {name}") + else: + pass + + torch.cuda.empty_cache() + + # save weights + if self.is_driver_worker: + from safetensors.torch import save_file + + tensors = {} + saved = False + count = 0 + size_in_bytes = 0 + + for name, weight in model.named_parameters(): + if "lm_head" in name and model.config.tie_word_embeddings: + continue + size_in_bytes += weight.numel() * weight.element_size() + tensors[name] = weight + saved = False + if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024: + weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6))) + save_file(tensors, weight_path) + print(f"The quantified weights were successfully saved in {weight_path}.") + tensors.clear() + saved = True + count += 1 + size_in_bytes = 0 + + if not saved: + weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6))) + save_file(tensors, weight_path) + print(f"The quantified weights were successfully saved in {weight_path}.") diff --git a/ixformer_sdk/contrib/vllm_flash_attn/__init__.py b/ixformer_sdk/contrib/vllm_flash_attn/__init__.py new file mode 100644 index 00000000..96c19ecb --- /dev/null +++ b/ixformer_sdk/contrib/vllm_flash_attn/__init__.py @@ -0,0 +1,3 @@ +__version__ = "2.6.1" + +from .flash_attn_interface import * diff --git a/ixformer_sdk/contrib/vllm_flash_attn/flash_attn_interface.py b/ixformer_sdk/contrib/vllm_flash_attn/flash_attn_interface.py new file mode 100644 index 00000000..f4727aad --- /dev/null +++ b/ixformer_sdk/contrib/vllm_flash_attn/flash_attn_interface.py @@ -0,0 +1,1018 @@ +import math +from typing import List, Optional, Union + +import ixformer._C as ops +import torch +from ixformer.inference.functions import vllm_paged_attention + +from ixformer.core import config + +__all__ = [ + "flash_attn_varlen_func", + "flash_attn_with_kvcache", + "ref_flash_attn_varlen_func", + "ref_flash_attn_with_kvcache", + "flash_attn_with_cache_batch_idx", + "flash_attn_decode_with_cache_batch_idx", + "ref_flash_attn_with_cache_batch_idx", + "flash_attn_prefill_with_cache_batch_idx", + "merge_attn_states", + "ref_merge_attn_states", +] + + +def ixinfer_flash_attn_unpad_wrapper( + q, + k, + v, + out, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + window_left, + window_right, + softmax_scale, + softcap, + sqrt_alibi, + alibi_slopes, +): + ops.infer.ixinfer_flash_attn_unpad( + q, + k, + v, + out, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + False, # need_lse =False + softmax_scale, + sqrt_alibi, + alibi_slopes, + ) + + +if config.IXFORMER_UNPAD_ATTENTION_ALGO == "ixinfer": + ixinfer_flash_attn_unpad_op = ops.infer.ixinfer_flash_attn_unpad_new +else: + ixinfer_flash_attn_unpad_op = ixinfer_flash_attn_unpad_wrapper + + +# https://github.com/vllm-project/flash-attention/blob/v2.6.2/vllm_flash_attn/flash_attn_interface.py +def flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, + block_table=None, + sqrt_alibi=False, + return_softmax_lse=False, + *, + out=None, +): + """dropout_p should be set to 0.0 during evaluation + Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Arguments: + q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch. + k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. + v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch. + cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths + of the sequences in the batch, used to index into q. + cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths + of the sequences in the batch, used to index into kv. + max_seqlen_q: int. Maximum query sequence length in the batch. + max_seqlen_k: int. Maximum key sequence length in the batch. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + deterministic: bool. Whether to use the deterministic implementation of the backward pass, + which is slightly slower and uses more memory. The forward pass is always deterministic. + return_attn_probs: bool. Whether to return the attention probabilities. This option is for + testing only. The returned probabilities are not guaranteed to be correct + (they might not have the right scaling). + Return: + out: (total, nheads, headdim). + softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen). + The output of softmax (possibly with different scaling). It also encodes the dropout + pattern (negative means that location was dropped, nonnegative means it was kept). + """ + + assert ( + deterministic is False + ), "For the inference model, we don't need this parameter." + assert ( + return_attn_probs is False + ), "For the inference model, we don't need this parameter." + assert dropout_p == 0, "For the inference model, we don't need this parameter." + + if out is None: + out = torch.empty_like(q) + + if softmax_scale is None: + softmax_scale = 1.0 / (q.size(-1) ** 0.5) + + num_tokens, head_num, head_dim = q.shape + lse = ( + torch.empty([head_num, num_tokens], device=q.device, dtype=torch.float32) + if return_softmax_lse + else None + ) + + if block_table is None: + ixinfer_flash_attn_unpad_op( + q, + k, + v, + out, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + window_size[0], + window_size[1], + softmax_scale, + softcap, + sqrt_alibi, + alibi_slopes, + lse, + ) + else: + ops.infer.ixinfer_flash_attn_unpad_with_block_tables( + q, + k, + v, + out, + block_table, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + window_size[0], + window_size[1], + softmax_scale, + softcap, + sqrt_alibi, + alibi_slopes, + lse, + ) + if return_softmax_lse: + return out, lse + return out + + +# https://github.com/vllm-project/flash-attention/blob/v2.6.1/vllm_flash_attn/flash_attn_interface.py#L1175 +def flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + block_table: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + alibi_slopes=None, + num_splits=0, + return_softmax_lse=False, + max_context_len: int = None, + use_cuda_graph: bool = False, + use_sqrt_alibi: bool = False, + *, + out=None, +): + """ + If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from + k and v. This is useful for incremental decoding: you can pass in the cached keys/values from + the previous step, and update them with the new keys/values from the current step, and do + attention with the updated cache, all in 1 kernel. + + If you pass in k / v, you must make sure that the cache is large enough to hold the new values. + For example, the KV cache could be pre-allocated with the max sequence length, and you can use + cache_seqlens to keep track of the current sequence lengths of each sequence in the batch. + + Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be + rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos + and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc. + If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at + indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens). + + See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function. + + Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV. + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V. + + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix. + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is: + 1 1 1 1 0 + 1 1 1 1 1 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is: + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + If the row of the mask is all zero, the output will be zero. + + If window_size != (-1, -1), implements sliding window local attention. Query at position i + will only attend to keys between + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. + + Note: Does not support backward pass. + + Arguments: + q: (batch_size, seqlen, nheads, headdim) + k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table, + or (num_blocks, nheads_k, page_block_size, headdim) if there's a block_table (i.e. paged KV cache) + page_block_size must be a multiple of 256. + v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table, + or (num_blocks, nheads_k, page_block_size, headdim) if there's a block_table (i.e. paged KV cache) + k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate + k with k_cache, starting at the indices specified by cache_seqlens. + v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k. + rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding + to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16. + rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos. + cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the + KV cache. + block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32. + cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache. + If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1]. + If the indices are not distinct, and k and v are provided, the values updated in the cache + might come from any of the duplicate indices. + softmax_scale: float. The scaling of QK^T before applying softmax. + Default to 1 / sqrt(headdim). + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling). + window_size: (left, right). If not (-1, -1), implements sliding window local attention. + softcap: float. Anything > 0 activates softcapping attention. + rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in. + If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False, + rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1 + (i.e. GPT-NeoX style). + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of + (-alibi_slope * |i + seqlen_k - seqlen_q - j|) + is added to the attention score of query i and key j. + num_splits: int. If > 1, split the key/value into this many chunks along the sequence. + If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic + to automatically determine the number of splits. + Don't change this unless you know what you are doing. + return_softmax_lse: bool. Whether to return the logsumexp of the attention scores. + + Return: + out: (batch_size, seqlen, nheads, headdim). + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax + normalization factor). + """ + assert k is None, "Updated *inplace* with the new key/values not supported." + assert v is None, "Updated *inplace* with the new key/values not supported." + assert rotary_cos is None and rotary_sin is None and cache_batch_idx is None + assert num_splits == 0 + assert return_softmax_lse is False + assert rotary_interleaved is True + + assert ( + max_context_len is not None + ), "flash_attn_with_kvcache needs to pass in the parameter 'max_context_len'." + + if out is None: + output = torch.empty_like(q) + else: + output = out + output_shape = list(output.shape) + + # For the official interface, the data layout is as follows： + # q: (batch_size, seqlen, nheads, headdim) + # k_cache, v_cache (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table + + # However, we adopts another data layout： + # q: (num_tokens, nheads, headdim) + # k_cache, v_cache (num_blocks, nheads_k, 16, headdim) if there's a block_table + batch_size, seqlen, nheads, headdim = q.shape + + # We assume shape is [num_blocks, nheads_k, page_block_size, headdim] + num_blocks, nheads_k, page_block_size, headdim = k_cache.shape + + assert page_block_size == 16 + q = q.view(batch_size * seqlen, nheads, headdim) + output = output.view(batch_size * seqlen, nheads, headdim) + + vllm_paged_attention( + output, + q, + k_cache, + v_cache, + nheads_k, + softmax_scale, + block_table, + cache_seqlens, + 16, + max_context_len, + alibi_slopes, + softcap, + True, + window_size[0], + window_size[1], + use_cuda_graph, + use_sqrt_alibi, + ) + + return output.view(*output_shape) + + +def flash_attn_prefill_with_cache_batch_idx( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_seqlens: torch.Tensor, + cache_batch_idx: torch.Tensor, + max_context_len: int, + softmax_scale: Optional[float] = None, + causal: Optional[bool] = False, + window_size: Optional[tuple] = (-1, -1), # -1 means infinite context window + softcap: Optional[float] = 0.0, # 0.0 means deactivated + alibi_slopes: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, +): + if output is None: + output = torch.empty_like(q) + + if softmax_scale is None: + head_dim = k_cache.shape[-1] + softmax_scale = 1 / head_dim**0.5 + + ops.infer.flash_attn_with_cache_batch_idx( + q, + k_cache, + v_cache, + output, + cache_seqlens, + cache_batch_idx, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + alibi_slopes, + ) + + return output + + +def flash_attn_decode_with_cache_batch_idx( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_seqlens: torch.Tensor, + cache_batch_idx: torch.Tensor, + max_context_len: int, + softmax_scale: float, + causal: bool, + window_size=(-1, -1), + softcap: float = 0, + alibi_slopes: torch.Tensor = None, + output: torch.Tensor = None, +): + if output is None: + output = torch.empty_like(q) + assert q.shape[1] == 1 + # q = q.view(q.shape[0], -1) + ops.infer.flash_attn_decode_with_cache_batch_idx( + q, + k_cache, + v_cache, + output, + cache_seqlens, + cache_batch_idx, + max_context_len, + softmax_scale, + causal, + window_size[0], + window_size[1], + softcap, + alibi_slopes, + ) + return output + + +def flash_attn_with_cache_batch_idx( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_seqlens: torch.Tensor, + cache_batch_idx: torch.Tensor, + max_context_len: int, + softmax_scale: Optional[float] = None, + causal: Optional[bool] = False, + window_size: Optional[tuple] = (-1, -1), # -1 means infinite context window + softcap: Optional[float] = 0.0, # 0.0 means deactivated + alibi_slopes: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, +): + """ + Args: + q: (batch_size, seqlen, nheads, headdim) torch.float16, torch.bfloat16 + k_cache: (batch_size_cache, nheads_k, seqlen_cache, headdim) torch.float16, torch.bfloat16 + v_cache: (batch_size_cache, nheads_k, seqlen_cache, headdim) torch.float16, torch.bfloat16 + cache_seqlens: (batch_size,) torch.int32 + cache_batch_idx: (batch_size,) torch.int32 + max_context_len: int + softmax_scale: float + causal: bool + window_size: tuple not implemented yet. + softcap: float not implemented yet. + alibi_slopes: (nheads,) torch.float32 causal must be true when alibi is not None + output: (batch_size, seqlen, nheads, headdim) torch.float16, torch.bfloat16 + Returns: + output: (batch_size, seqlen, nheads, headdim) torch.float16, torch.bfloat16 + """ + assert len(q.shape) == 4 + if ( + q.shape[1] == 1 and window_size[0] == -1 and window_size[1] == -1 + ): # remove window size check when kernel supported. + return flash_attn_decode_with_cache_batch_idx( + q=q, + k_cache=k_cache, + v_cache=v_cache, + cache_seqlens=cache_seqlens, + cache_batch_idx=cache_batch_idx, + max_context_len=max_context_len, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, + output=output, + ) + else: + return flash_attn_prefill_with_cache_batch_idx( + q=q, + k_cache=k_cache, + v_cache=v_cache, + cache_seqlens=cache_seqlens, + cache_batch_idx=cache_batch_idx, + max_context_len=max_context_len, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + softcap=softcap, + alibi_slopes=alibi_slopes, + output=output, + ) + + +def get_alibi_mask(num_heads, seqlen, device, dtype, sqrt_alibi=False): + offsets = torch.arange(seqlen) + offsets = offsets[None, :] - offsets[:, None] + offsets = offsets.to(device) + if sqrt_alibi: # sqrt distance for alibi bias + offsets = torch.sqrt(torch.abs(offsets)) * torch.sign(offsets) + + return offsets + + +def get_alibi_mask_decode(num_heads, seqlen, device, dtype, sqrt_alibi=False): + x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1) + y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1) + if sqrt_alibi: + offsets = -torch.sqrt((y - x)).view(1, 1, seqlen) + else: + offsets = -(y - x).view(1, 1, seqlen) + return offsets + + +def construct_local_mask( + seqlen_q, + seqlen_k, + window_size=(-1, -1), # -1 means infinite window size + device=None, +): + row_idx = torch.arange(seqlen_q, device=device, dtype=torch.long).view(-1, 1) + mask = ( + torch.arange(seqlen_k, device=device, dtype=torch.long) + .view(1, -1) + .repeat(seqlen_q, 1) + ) + # [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] + return ((row_idx + seqlen_k - seqlen_q - window_size[0]) <= mask) & ( + (row_idx + seqlen_k - seqlen_q + window_size[1]) >= mask + ) + + +def compute_softmax_lse(x): + # 为了数值稳定性，先减去最大值 + input_tensor = x + if x.shape[-1] == 0: + lse = torch.full(x.shape[:-1], float("inf"), device=x.device) + softmax_out = torch.empty(x.shape, dtype=x.dtype, device=x.device) + return softmax_out, lse.view(x.shape[:-1]) + max_values = torch.max(input_tensor, dim=-1, keepdim=True)[0] + input_tensor = input_tensor - max_values + + # 计算以 2 为底的指数 + log2_e = math.log2(math.e) + exp2_tensor = torch.exp2(input_tensor * log2_e) + + # 计算指数和 + exp2_sum = torch.sum(exp2_tensor, dim=-1, keepdim=True) + + # 计算 softmax + softmax_output = exp2_tensor / exp2_sum + + lse = max_values * log2_e + torch.log2(exp2_sum) + return softmax_output, lse.view(x.shape[:-1]) + + +def ref_flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, + block_table=None, + sqrt_alibi=False, + return_softmax_lse=False, + *, + out=None, +) -> torch.Tensor: + num_seqs = len(cu_seqlens_q) - 1 + num_tokens, num_query_heads, head_size = q.shape + if block_table is None: + _, num_kv_heads, _ = k.shape + else: + _, num_kv_heads, _, _ = k.shape + + num_query_heads = q.shape[1] + slopes = ( + alibi_slopes.view(num_query_heads, 1, 1) + if alibi_slopes is not None + else alibi_slopes + ) + + if softmax_scale is None: + softmax_scale = 1.0 / (q.size(-1) ** 0.5) + + outputs: List[torch.Tensor] = [] + + # head_num, num_tokens + lse = torch.empty([q.shape[1], q.shape[0]], dtype=torch.float32, device=q.device) + + for i in range(num_seqs): + query_start_idx = cu_seqlens_q[i] + query_end_idx = cu_seqlens_q[i + 1] + kv_start_idx = cu_seqlens_k[i] + kv_end_idx = cu_seqlens_k[i + 1] + + query_len = query_end_idx - query_start_idx + kv_len = kv_end_idx - kv_start_idx + + sq = q[query_start_idx:query_end_idx] + + if block_table is None: + sk = k[kv_start_idx:kv_end_idx] + sv = v[kv_start_idx:kv_end_idx] + else: + table = block_table[i] + ks = [] + vs = [] + need_blocks = (kv_len + 15) // 16 + for index in range(need_blocks): + offset = ( + 16 + if index != (need_blocks - 1) + else min(16, 16 if kv_len % 16 == 0 else kv_len % 16) + ) + ks.append(k[table[index], :, :offset, :].permute(1, 0, 2)) + vs.append(v[table[index], :, :offset, :].permute(1, 0, 2)) + sk = torch.cat(ks, dim=0) + sv = torch.cat(vs, dim=0) + assert sk.shape[0] == kv_len + assert sv.shape[0] == kv_len + + if num_query_heads != num_kv_heads: + assert ( + num_query_heads > num_kv_heads and num_query_heads % num_kv_heads == 0 + ) + sk = torch.repeat_interleave(sk, num_query_heads // num_kv_heads, dim=1) + sv = torch.repeat_interleave(sv, num_query_heads // num_kv_heads, dim=1) + attn = torch.einsum("qhd,khd->hqk", sq, sk * softmax_scale).float() + # 0 -> mask out 1 -> calculation + mask = torch.ones(query_len, kv_len, device=q.device) + shift = ( + kv_len - query_len + ) # flash attention use bottom-right as default, so we do not use shift = 0 + if causal: + mask = torch.tril(mask, diagonal=shift).bool() + else: + mask = mask.bool() + + if window_size[0] != -1 and window_size[1] != -1: + # [left, right] + win_mask = construct_local_mask( + query_len, + kv_len, + window_size=( + window_size[0], + window_size[1], + ), # -1 means infinite window size + device=mask.device, + ) + mask = win_mask & mask + elif window_size[1] != -1: + # [-1, right] + right = window_size[1] + win_mask = torch.tril(mask, diagonal=shift + right).bool() + mask = win_mask & mask + elif window_size[0] != -1: + # [left, -1] + left = window_size[0] + win_mask = torch.triu(mask, diagonal=shift - left).bool() + mask = win_mask & mask + + # 1 -> 0, 0 -> 1 to mask out + zero_index = ~(mask.sum(dim=-1).bool()) + mask = ~mask + mask = mask.float() * -10000 + + if alibi_slopes is not None: + offsets = get_alibi_mask( + num_query_heads, kv_len, q.device, q.dtype, sqrt_alibi + ) + alibi_mask = offsets * slopes + alibi_mask = alibi_mask.to(attn.dtype) + alibi_mask = alibi_mask[:, -query_len:] + + attn = attn + mask.to(attn.dtype).to(attn.device) # num_heads, kv_len, head_dim + + if alibi_slopes is not None: + attn = attn + alibi_mask + + # attn.masked_fill_(mask, float("-inf")) + # attn = torch.softmax(attn, dim=-1).to(sv.dtype) + attn, tmp_lse = compute_softmax_lse(attn) + attn = attn.to(sv.dtype) + + sout = torch.einsum("hqk,khd->qhd", attn, sv) + if softcap != 0: + sout = softcap * torch.tanh(sout / softcap) + sout[zero_index] = 0.0 + outputs.append(sout) + + lse[:, query_start_idx:query_end_idx] = tmp_lse + + outputs = torch.cat(outputs, dim=0) + + if out is not None: + out.copy_(outputs) + else: + out = outputs + if return_softmax_lse: + return out, lse + return out + + +def ref_flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + rotary_cos=None, + rotary_sin=None, + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + block_table: Optional[torch.Tensor] = None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + softcap=0.0, # 0.0 means deactivated + rotary_interleaved=True, + alibi_slopes=None, + num_splits=0, + return_softmax_lse=False, + max_context_len: int = None, + use_sqrt_alibi: bool = False, + *, + out=None, +) -> torch.Tensor: + assert k is None + assert v is None + assert rotary_cos is None + assert rotary_sin is None + assert cache_batch_idx is None + assert causal is True + assert rotary_interleaved + assert num_splits == 0 + assert not return_softmax_lse + + head_size = q.size(-1) + + num_seqs = cache_seqlens.size(0) + block_tables = block_table.cpu().numpy() + + _, num_kv_heads, block_size, head_size = k_cache.shape + + assert block_size == 16 + + num_query_heads = q.shape[-2] + q_shape = q.shape + q = q.view(-1, num_query_heads, head_size) + + slopes = ( + alibi_slopes.view(num_query_heads, 1, 1) + if alibi_slopes is not None + else alibi_slopes + ) + + outputs: List[torch.Tensor] = [] + + for i in range(num_seqs): + kv_len = cache_seqlens[i].item() + + sq = q[i : i + 1] + + num_kv_blocks = (kv_len + block_size - 1) // block_size + block_indices = block_tables[i, :num_kv_blocks] + + sk = k_cache[block_indices].permute( + 0, 2, 1, 3 + ) # -> num_blocks, block_size, num_head, head_size + sk = sk.reshape(-1, num_kv_heads, head_size) + sk = sk[:kv_len] + sv = v_cache[block_indices].permute(0, 2, 1, 3) + sv = sv.reshape(-1, num_kv_heads, head_size) + sv = sv[:kv_len] + + if num_query_heads != num_kv_heads: + sk = torch.repeat_interleave(sk, num_query_heads // num_kv_heads, dim=1) + sv = torch.repeat_interleave(sv, num_query_heads // num_kv_heads, dim=1) + + attn = torch.einsum("qhd,khd->hqk", sq, sk * softmax_scale).float() + empty_mask = torch.ones(1, kv_len) + mask = torch.triu(empty_mask, diagonal=kv_len).bool().to(q.device) + + if window_size != (-1, -1): + # sliding_window_mask = torch.triu(empty_mask, + # diagonal=kv_len - + # (query_len + sliding_window) + + # 1).bool().logical_not() + sliding_window_mask = construct_local_mask( + 1, + kv_len, + window_size=( + window_size[0], + window_size[1], + ), # -1 means infinite window size + device=mask.device, + ) + mask |= sliding_window_mask + mask = mask.float() * -1000 + + if alibi_slopes is not None: + offsets = get_alibi_mask_decode( + num_query_heads, kv_len, q.device, q.dtype, use_sqrt_alibi + ) + alibi_mask = offsets * slopes + alibi_mask = alibi_mask.to(attn.dtype) + + attn = attn + mask.to(attn.dtype) # num_heads, kv_len, head_dim + + if alibi_slopes is not None: + attn = attn + alibi_mask + + # attn.masked_fill_(mask, float("-inf")) + attn = torch.softmax(attn, dim=-1).to(sv.dtype) + sout = torch.einsum("hqk,khd->qhd", attn, sv) + if softcap != 0: + sout = softcap * torch.tanh(sout / softcap) + outputs.append(sout) + + outputs = torch.cat(outputs, dim=0) + + if out is not None: + out_shape = out.shape + out = out.view(*outputs.shape) + out.copy_(outputs) + else: + out_shape = q_shape + out = outputs + + return out.view(*out_shape) + + +def ref_flash_attn_with_cache_batch_idx( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_seqlens: torch.Tensor, + cache_batch_idx: torch.Tensor, + max_context_len: int, + softmax_scale: Optional[float] = None, + causal: Optional[bool] = False, + window_size: Optional[tuple] = (-1, -1), # -1 means infinite context window + softcap: Optional[float] = 0.0, # 0.0 means deactivated + alibi_slopes: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, +): + basic_output = torch.empty_like(q) + cache_batch_idx_cpu = cache_batch_idx.cpu() + cache_seqlens_cpu = cache_seqlens.cpu() + + for i, batch_idx in enumerate(cache_batch_idx_cpu): + cur_q = q[i] + cur_k = k_cache[batch_idx, :, : cache_seqlens_cpu[i], :] + cur_v = v_cache[batch_idx, :, : cache_seqlens_cpu[i], :] + + cur_q = cur_q.transpose(0, 1) + # cur_k = cur_k.transpose(0, 1) + # cur_v = cur_v.transpose(0, 1) + + num_q_heads = cur_q.shape[0] + num_kv_heads = cur_k.shape[0] + + assert num_q_heads >= num_kv_heads and num_q_heads % num_kv_heads == 0 + kv_repeat = num_q_heads // num_kv_heads + if kv_repeat > 1: + cur_k = cur_k.repeat_interleave(kv_repeat, dim=0) + cur_v = cur_v.repeat_interleave(kv_repeat, dim=0) + attn = torch.matmul(cur_q, cur_k.transpose(-1, -2)).float() * softmax_scale + + query_len = cur_q.shape[1] + kv_len = cur_k.shape[1] + + mask = None + if causal: + mask = torch.ones(query_len, kv_len, device=q.device) + shift = kv_len - query_len + mask = torch.tril(mask, diagonal=shift).bool() + + if window_size[0] != -1 and window_size[1] != -1: + # [left, right] + win_mask = construct_local_mask( + query_len, + kv_len, + window_size=( + window_size[0], + window_size[1], + ), # -1 means infinite window size + device=mask.device, + ) + mask = win_mask & mask + elif window_size[1] != -1: + # [-1, right] + right = window_size[1] + win_mask = torch.tril(mask, diagonal=shift + right).bool() + mask = win_mask & mask + elif window_size[0] != -1: + # [left, -1] + left = window_size[0] + win_mask = torch.triu(mask, diagonal=shift - left).bool() + mask = win_mask & mask + + if mask is not None: + attn.masked_fill_(~mask, float("-inf")) + + if alibi_slopes is not None: + slopes = alibi_slopes.view(num_q_heads, 1, 1) + offsets = get_alibi_mask_decode(num_q_heads, kv_len, q.device, q.dtype) + alibi_mask = offsets * slopes + alibi_mask = alibi_mask.to(attn.dtype) + attn = attn + alibi_mask + + attn = torch.softmax(attn, dim=-1).to(q.dtype) + if mask is not None: + attn.masked_fill_(~mask, 0) + + out = torch.matmul(attn, cur_v) + basic_output[i, :, :, :] = out.transpose(0, 1) + if output is None: + output = basic_output + else: + output.copy_(basic_output) + return output + + +def ref_merge_attn_states( + out_1, lse_1, out_2, lse_2, output=None, return_lse: Optional[bool] = False +): + lse_1 = torch.where( + lse_1 == float("inf"), torch.full_like(lse_1, -float("inf")), lse_1 + ) + lse_2 = torch.where( + lse_2 == float("inf"), torch.full_like(lse_2, -float("inf")), lse_2 + ) + num_heads, seq_len = lse_1.shape + + lse_2 = lse_2.transpose(0, 1).view(seq_len, num_heads, 1) + lse_1 = lse_1.transpose(0, 1).view(seq_len, num_heads, 1) + + s_max = torch.maximum(lse_1, lse_2) + + d = torch.exp2(lse_1 - s_max) + torch.exp2(lse_2 - s_max) + v_merged = out_1 * torch.exp2(lse_1 - s_max) + out_2 * torch.exp2(lse_2 - s_max) + v_merged = v_merged / d + v_merged = v_merged.to(out_1.dtype) + if output is None: + output = v_merged + else: + output.copy_(v_merged) + if return_lse: + output_lse = s_max + torch.log2(d) + output_lse = output_lse.squeeze(-1).transpose(0, 1).contiguous() + return output, output_lse + else: + return output + + +def merge_attn_states( + prefix_output: torch.Tensor, + prefix_lse: torch.Tensor, + suffix_output: torch.Tensor, + suffix_lse: torch.Tensor, + output: torch.Tensor = None, + return_lse: Optional[bool] = False, +): + """ + Args: + prefix_output: (seq_len, head_num, head_dim) torch.float16, torch.bfloat16 + prefix_lse: (head_num, seq_len) torch.float32 + suffix_output: (seq_len, head_num, head_dim) torch.float16, torch.bfloat16 + suffix_lse: (head_num, seq_len) torch.float32 + Returns: + output: (seq_len, head_num, head_dim) torch.float16, torch.bfloat16 + output_lse: (head_num, seq_len) torch.float32 + """ + if output is None: + output = torch.empty_like(prefix_output) + output_lse = torch.empty_like(prefix_lse) if return_lse else None + + ops.infer.merge_attn_states( + prefix_output, prefix_lse, suffix_output, suffix_lse, output, output_lse + ) + if output_lse is not None: + return output, output_lse + else: + return output diff --git a/ixformer_sdk/core/__init__.py b/ixformer_sdk/core/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/core/config.py b/ixformer_sdk/core/config.py new file mode 100644 index 00000000..039db399 --- /dev/null +++ b/ixformer_sdk/core/config.py @@ -0,0 +1,184 @@ +import os +from typing import Callable, Optional + +# ========================================================= +# Utils +# ========================================================= + + +def number_type(scalar_type): + def wrap(val: Optional[str]): + if val is None: + return None + + return scalar_type(val) + + return wrap + + +def bool_type(val: Optional[str]): + if val is None: + return False + + if isinstance(val, str): + return val.lower() in ["1", "t", "true"] + + if isinstance(val, int): + return val != 0 + + raise RuntimeError(f"Invalid bool type, got {type(val), val}") + + +def list_type(scalar_type=str): + def wrap(val: Optional[str]): + if val is None: + return [] + + if not isinstance(val, str): + raise RuntimeError( + f"list_type: Got invalid type, expect str, but got {val}." + ) + + return [scalar_type(v) for v in val.split(",")] + + return wrap + + +def Field( + name: str, + static: bool = True, + type: Callable = str, + choices: Optional[list] = None, + help: Optional[str] = None, + **kwargs, +): + """ + Define environment variable field + + Example: + Static mode: + # define + ENABLE_XX = Field("ENABLE_XX", type=bool, help="ENABLE_XX") + + # use + config.ENABLE_XX + + Dynamic mode: + # Please use lowercase naming to differentiate it with static mode. + + # define + enable_cc = Field("ENABLE_CC", type=bool, static=False, help="enable_cc") + + # use + config.enable_cc() + + Set default value: + # define + ENABLE_TT = Field("ENABLE_TT", type=bool, default=False, help="ENABLE_TT") + + # use + config.ENABLE_TT + + Use list: + # define + CUDA_VISIBLE_DEVICES = Field("CUDA_VISIBLE_DEVICES", type=list_type(int), help="CUDA_VISIBLE_DEVICES") + + # use + # the CUDA_VISIBLE_DEVICES is parsed to list, and it's value is int type. + for device_id in CUDA_VISIBLE_DEVICES: + ... + + """ + + if type == bool: + type = bool_type + + elif type in [list, tuple]: + type = list_type(scalar_type=str) + + elif type in [int, float]: + type = number_type(type) + + if static: + env_val = type(os.environ.get(name, **kwargs)) + if choices is not None and env_val is not None and env_val not in choices: + raise RuntimeError( + f"Got invalid value, expect {choices}, but got {env_val}." + ) + return env_val + + def _get(): + env_val = type(os.environ.get(name, **kwargs)) + if choices is not None and env_val is not None and env_val not in choices: + raise RuntimeError( + f"Got invalid value, expect {choices}, but got {env_val}." + ) + return env_val + + return _get + + +# ========================================================= +# Functions Config +# ========================================================= + +IXFORMER_GEMV_THRESHOLD = Field( + "IXFORMER_GEMV_THRESHOLD", + type=int, + default=1, + help="Set the threshold for using gemv.", +) + + +# ========================================================= +# Distributed Config +# ========================================================= + +IXFORMER_COMM_SHM_SIZE = Field( + "IXFORMER_COMM_SHM_SIZE", + type=int, + default=None, + help="set shared memory size of ipc comm.", +) + +IXFORMER_ENABLE_OVERLAP_COMM = Field( + "IXFORMER_ENABLE_OVERLAP_COMM", + type=bool, + default=False, + help="enable overlap communcation and compute.", +) + +IXFORMER_OVERLAP_GEMM_METHOD = Field( + "IXFORMER_OVERLAP_GEMM_METHOD", + type=int, + default=None, + choices=[0, 1], + help="set gemm backend, 0: ixinfer, 1: cublas.", +) + +IXFORMER_OVERLAP_CHUNKS = Field( + "IXFORMER_OVERLAP_CHUNKS", type=int, default=2, help="set split chunks." +) + +IXFORMER_OVERLAP_SPLIT_RATIO = Field( + "IXFORMER_OVERLAP_SPLIT_RATIO", + type=float, + default=None, + help="set split chunks ratio.", +) + +IXFORMER_PAGED_ATTENTION_ALGO = Field( + "IXFORMER_PAGED_ATTENTION_ALGO", + type=str, + default="ixinfer", + choices=["ixinfer", "ixformer"], + help="set paged attention algo.", +) + +IXFORMER_UNPAD_ATTENTION_ALGO = Field( + "IXFORMER_UNPAD_ATTENTION_ALGO", + type=str, + default="ixinfer", + choices=["ixinfer", "ixinfer-ex"], + help="set enpad attention algo.", +) diff --git a/ixformer_sdk/core/dispatcher.py b/ixformer_sdk/core/dispatcher.py new file mode 100644 index 00000000..57fb91b1 --- /dev/null +++ b/ixformer_sdk/core/dispatcher.py @@ -0,0 +1,20 @@ +class Dispatcher(object): + """ + create object by dispatcher to reuse object. + """ + + _dispatcher = dict() + + @classmethod + def dispatcher(cls, *args, **kwargs): + key = cls.dispatcher_key(*args, **kwargs) + obj = cls._dispatcher.get(key, None) + if obj is None: + obj = cls(*args, **kwargs) + cls._dispatcher[key] = obj + + return obj + + @classmethod + def dispatcher_key(cls, *args, **kwargs): + raise NotImplementedError() diff --git a/ixformer_sdk/core/multi_level_cache.py b/ixformer_sdk/core/multi_level_cache.py new file mode 100644 index 00000000..1e42d2b5 --- /dev/null +++ b/ixformer_sdk/core/multi_level_cache.py @@ -0,0 +1,54 @@ +class MultiLevelCache(object): + def __init__(self): + self._l1_key = None + self._l1_value = None + + self._l2_size = 3 + self._l2 = [(None, None) for _ in range(self._l2_size)] + self._l2_ptr = 0 + + self._l3 = dict() + + def set(self, key, value): + self._l1_key = key + self._l1_value = value + + self._l2[self._l2_ptr] = (key, value) + self._l2_ptr = (self._l2_ptr + 1) % 3 # l2_size: 3 + + self._l3[key] = value + + def get(self, key, *args): + if key == self._l1_key: + return self._l1_value + + l2 = self._l2 + if key == l2[0][0]: + return l2[0][1] + + if key == l2[1][0]: + return l2[1][1] + + if key == l2[2][0]: + return l2[2][1] + + return self._l3.get(key, *args) + + def containe(self, key): + return key in self._l3 + + def __getitem__(self, item): + return self.get(item) + + def __setitem__(self, key, value): + self.set(key, value) + + def __contains__(self, item): + if item == self._l1_key: + return True + + l2 = self._l2 + if item == l2[0][0] or item == l2[1][0] or item == l2[2][0]: + return True + + return item in self._l3 diff --git a/ixformer_sdk/core/operator_autotuning.py b/ixformer_sdk/core/operator_autotuning.py new file mode 100644 index 00000000..84a07d57 --- /dev/null +++ b/ixformer_sdk/core/operator_autotuning.py @@ -0,0 +1,237 @@ +import abc +import bisect +import functools +import itertools +import random +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union + +import torch +import torch.distributed as dist + +import ixformer.distributed as ixfd +from ixformer.utils.benchmark.cuda_benchmark import Functor, cuda_benchmark + + +def sync_ranks_metric(value, group=None): + if not isinstance(value, (torch.Tensor, int, float)): + raise RuntimeError( + f"Invalid metric value, expect `Tensor`, `int`, or `float` type, but got {value}." + ) + + if torch.is_tensor(value): + value = value.to("cuda") + else: + value = torch.tensor([value], dtype=torch.float, device="cuda") + + dist.broadcast(value, src=0, group=group) + return value.cpu().item() + + +class AutotuningFinder(object): + def freeze(self): + pass + + @abc.abstractmethod + def get(self, key) -> Callable: + pass + + @abc.abstractmethod + def set(self, *args, **kwargs): + pass + + +class BasedKeyFinder(AutotuningFinder): + def __init__(self): + self._key_to_value: Dict[Any, Callable] = dict() + + def get(self, key, **kwargs) -> Callable: + if "default" in kwargs: + return self._key_to_value.get(key, kwargs["default"]) + return self._key_to_value[key] + + def set(self, key, value): + self._key_to_value[key] = value + + def containe(self, key): + return key in self._key_to_value + + +class TreeNode: + def __init__(self): + self.nodes: List[Union[Any, TreeNode]] = list() + self.key_to_nodes: Dict[Any, TreeNode] = dict() + + def add(self, key, value): + if isinstance(key, (tuple, list)): + if len(key) == 1: + self.insert_value(key[0], value) + else: + self.recurse_add_node(key, value) + else: + self.insert_value(key, value) + + def insert_value(self, key, value): + self.nodes.append((key, value)) + self.key_to_nodes[key] = value + + def recurse_add_node(self, key, value): + if key[0] in self.key_to_nodes: + node = self.key_to_nodes[key[0]] + else: + node = TreeNode() + self.key_to_nodes[key[0]] = node + self.insert_value(key[0], node) + + node.add(key[1:], value) + + def sort(self): + self.nodes.sort(key=lambda x: x[0]) + for _, node in self.nodes: + if isinstance(node, TreeNode): + node.sort() + + def find(self, key): + is_list_key = isinstance(key, (tuple, list)) + if not is_list_key: + key = (key,) + + num_querys = len(key) + node = self + for key_idx in range(num_querys): + query_key = key[key_idx] + idx = bisect.bisect_left(node.nodes, (query_key,)) - 1 + if idx <= 0: + node = node.nodes[0][1] + elif idx >= len(node.nodes): + node = node.nodes[-1][1] + else: + node = node.nodes[idx][1] + + return node + + def show(self, indent=0): + for k, node in self.nodes: + print(" " * indent, end="") + if isinstance(node, TreeNode): + print(f"key: {k}") + node.show(indent=indent + 4) + else: + print(f"key: {k}, node: {node}") + + +class BasedRangeFinder(AutotuningFinder): + def __init__(self): + super().__init__() + + self.tree = TreeNode() + self._found_cache: Dict[Any, Callable] = dict() + + def freeze(self): + self.tree.sort() + + def get(self, key) -> Callable: + value = self._found_cache.get(key, None) + if value is not None: + return value + + value = self.tree.find(key) + self._found_cache[key] = value + return value + + def set(self, key, value): + self.tree.add(key, value) + + +class OperatorAutotuning(object): + def __init__(self, num_repeated=5, num_warmup=3, dist_barrier=False): + self.num_repeated = num_repeated + self.num_warmup = num_warmup + self.dist_barrier = dist_barrier + + @abc.abstractmethod + def operators(self): + raise NotImplementedError() + + def __call__(self, *args, **kwargs): + return self.exec_best_operator(args, kwargs) + + @abc.abstractmethod + def exec_best_operator(self, args, kwargs): + raise NotImplementedError() + + @abc.abstractmethod + def autotuning(self, *args, **kwargs): + raise NotImplementedError() + + def perf_best_operator(self, *args, **kwargs) -> Callable: + best_operator = None + best_operator_time = float("inf") + + for idx, operator in enumerate(self.operators()): + op_time = self.perf_operator_time(operator, *args, **kwargs) + if op_time < best_operator_time: + best_operator = operator + best_operator_time = op_time + + # print(operator, op_time) + + return best_operator + + def perf_operator_time(self, op: Callable, *args, **kwargs) -> float: + fn = Functor(op, *args, **kwargs) + time = cuda_benchmark(fn, self.num_repeated, self.num_warmup, self.dist_barrier) + return time.gpu + + +class OperatorRuntimeAutotuning(OperatorAutotuning): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.operator_finder = BasedKeyFinder() + + @abc.abstractmethod + def get_operator_key(self, *args, **kwargs): + raise NotImplementedError() + + def exec_best_operator(self, args, kwargs): + key = self.get_operator_key(*args, **kwargs) + operator = self.operator_finder.get(key, default=None) + if operator is None: + operator = self.perf_best_operator(*args, **kwargs) + self.operator_finder.set(key, operator) + + return operator(*args, **kwargs) + + +class OperatorPreBaseRangeAutotuning(OperatorAutotuning): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.operator_finder = BasedRangeFinder() + self._finished_autotuning = False + + @abc.abstractmethod + def get_operator_key(self, *args, **kwargs): + raise NotImplementedError() + + @abc.abstractmethod + def generate_operator_inputs(self) -> Iterable[Tuple[Tuple, Dict]]: + raise NotImplementedError() + + def exec_best_operator(self, args, kwargs): + if not self._finished_autotuning: + self.autotuning() + + best_op = self.operator_finder.get(self.get_operator_key(*args, **kwargs)) + return best_op(*args, **kwargs) + + def autotuning(self): + for op_args, op_kwargs in self.generate_operator_inputs(): + best_op = self.perf_best_operator(*op_args, **op_kwargs) + self.operator_finder.set( + self.get_operator_key(*op_args, **op_kwargs), best_op + ) + + self.operator_finder.freeze() + self._finished_autotuning = True + # self.operator_finder.tree.show() diff --git a/ixformer_sdk/csrc/FindIXFORMER.cmake b/ixformer_sdk/csrc/FindIXFORMER.cmake new file mode 100644 index 00000000..42988085 --- /dev/null +++ b/ixformer_sdk/csrc/FindIXFORMER.cmake @@ -0,0 +1,40 @@ +# use python to find ixformer libs and include +if (CMAKE_VERSION VERSION_LESS 3.18) + set(DEV_MODULE Development) +else() + set(DEV_MODULE Development.Module) +endif() + +find_package(Python COMPONENTS Interpreter ${DEV_MODULE} REQUIRED) + +# find ixformer +set(IXFORMER_FOUND FALSE) + +if("${Python_FOUND}" STREQUAL "TRUE") + execute_process( + COMMAND ${Python_EXECUTABLE} -c "import os, ixformer; print(os.path.dirname(ixformer.__file__))" + OUTPUT_VARIABLE IXFORMER_PYDIR + ERROR_VARIABLE PYTHON_ERROR + RESULT_VARIABLE PYTHON_RESULT + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_STRIP_TRAILING_WHITESPACE + ) + if ("${IXFORMER_PYDIR}" STREQUAL "") + message("-- Not found ixFormer") + else () + message("-- Found ixFormer: ${IXFORMER_PYDIR}") + set(IXFORMER_FOUND TRUE) + endif () +endif() + +set(IXFORMER_COMM_LIBS "ixformer_comm") +set(IXFORMER_KERNEL_LIBS "ixformer_kernels") +set(IXFORMER_LIBS "${IXFORMER_COMM_LIBS} ${IXFORMER_KERNEL_LIBS}") +set(IXFORMER_INCLUDE "") +set(IXFORMER_DIR "") + +if("${IXFORMER_FOUND}" STREQUAL "TRUE") + set(IXFORMER_INCLUDE "${IXFORMER_PYDIR}/csrc/include") + set(IXFORMER_DIR "${IXFORMER_PYDIR}") + message("-- ixFormer LIBS: ${IXFORMER_LIBS}, INCLUDE: ${IXFORMER_INCLUDE}") +endif () diff --git a/ixformer_sdk/csrc/include/ixformer/comm/ccl.h b/ixformer_sdk/csrc/include/ixformer/comm/ccl.h new file mode 100644 index 00000000..f227002a --- /dev/null +++ b/ixformer_sdk/csrc/include/ixformer/comm/ccl.h @@ -0,0 +1,357 @@ +#pragma once + +#include "core/op_algo.h" +#include "nccl.h" + +namespace ixformer::comm { + +const uint8_t MAX_TENSOR_NDIM = 8; +constexpr size_t DEFAULT_SHM_SIZE = 16 * 1024 * 1024 * sizeof(float); + +struct Comm; +typedef Comm *Comm_t; + + +struct TensorDesc { + void *data_ptr; + ncclDataType_t dtype; + uint64_t numel; + uint8_t ndim; + int64_t shape[MAX_TENSOR_NDIM]; + int64_t stride[MAX_TENSOR_NDIM]; + bool contiguous; +}; + + +/** + * @brief Generate unique communicator id. + * + * Generates an Id to be used in ncclCommInitRank. ncclGetUniqueId should be + * called once and the Id should be distributed to all ranks in the + * communicator before calling ncclCommInitRank. + * + * @param commId: the unique id of communicator, it is created in main rank, and broadcast other rank. + */ +void getUniqueId(ncclUniqueId *commId); + +/** + * @brief Serialize commId to string. + * @param commId: the unique id of communicator. + * @return: serialized string. + */ +std::string serializeUniqueId(const ncclUniqueId &commId); + +/** + * @brief Deserialize the string of commId. + * @param commIdStr: serialized string by serializeUniqueId. + * @param commId: output commId + */ +void deserializeUniqueId(const std::string &commIdStr, ncclUniqueId *commId); + +/** + * @brief Creates a new communicator (multi process version). + * + * Rank must be between 0 and nranks-1 and unique within a communicator clique. + * Each rank is associated to a CUDA device, which has to be set before calling ncclCommInitRank. + * + * It is important to ensure that the current process's CUDA device is set by cudaSetDevice before calling this function, + * otherwise, an exception will be thrown. + * + * @param comm: Communicator + * @param nranks: the number of ranks. + * @param commId: the unique id of communicator. + * @param rank: the rank of current process + * @param shm_size: Unlike NCCL, IxFormer communication relies on CUDA IPC for communication by shared memory. + * If the shm_size is not provided, it will use the default value: DEFAULT_SHM_SIZE. + * @throw CommError: Throw CommError when an error is encountered. + */ +void initRank(Comm_t *comm, int nranks, ncclUniqueId commId, int rank, size_t shm_size = DEFAULT_SHM_SIZE); + +/** + * @brief Finalize a communicator. + * + * ncclCommFinalize flushes all issued communications, + * and marks communicator state as ncclInProgress. The state will change to ncclSuccess + * when the communicator is globally quiescent and related resources are freed; then, + * calling ncclCommDestroy can locally free the rest of the resources (e.g. communicator + * itself) without blocking. + * + * @param comm: Communicator + * @throw CommError: Throw CommError when an error is encountered. + */ +void destroy(Comm_t comm); + +void delete_comm_resuouces(Comm_t comm); + +/** + * @brief Whether is initiated. + * @param comm: Communicator + */ +bool isInitiated(Comm_t comm); + +/** + * @brief Gets ncclComm_t. + * @param comm: Communicator + */ +ncclComm_t getNcclComm(Comm_t comm); + +/** + * @brief Gets the number of ranks in the communicator clique + * @param comm: Communicator + */ +int getWorldSize(Comm_t comm); + +/** + * @brief Gets the number of nodes in the communicator clique + * @param comm: Communicator + */ +int getNumNodes(Comm_t comm); + +/** + * @brief Returns the user-ordered "rank" associated with the communicator. + * @param comm: Communicator + */ +int getRank(Comm_t comm); + +/** + * @brief Returns the cuda device number associated with the communicator. + * @param comm: Communicator + */ +int getDevice(Comm_t comm); + +/** + * @brief Gets shared memory size in the communicator clique + * @param comm: Communicator + */ +uint64_t getIpcShmSize(Comm_t comm); + + +// ============================================================================ +// Collective communication operations +// +// Collective communication operations must be called separately for each +// communicator in a communicator clique. +// +// They return when operations have been enqueued on the CUDA stream. +// +// Since they may perform inter-CPU synchronization, each call has to be done +// from a different thread or process, or need to use Group Semantics (see +// below). +// ============================================================================ + +/** + * @brief Barrier the member of the communicator + * @param comm: Communicator + * @param stream: CUDA Stream + */ +void barrier(Comm_t comm, cudaStream_t stream); + +/** + * @brief All-Gather + * + * Each device gathers sendcount values from other GPUs into senddata, + * receiving data from rank i at offset i*sendcount. + * Assumes recvcount is equal to nranks*sendcount, which means that recvdata + * should have a size of at least nranks*sendcount elements. + * + * In-place operations will happen if senddata == recvdata + rank * sendcount. + * + * @param comm: Communicator + * @param senddata: send data + * @param recvdata: recv data + * @param sendcount: the number of send elements, it is not nbytes. + * @param dtype: data type + * @param stream: CUDA Stream + * @param algo: Algorithm + * @throw CommError: Throw CommError when an error is encountered. + */ +void allGather(Comm_t comm, const void *senddata, void *recvdata, size_t sendcount, ncclDataType_t dtype, + cudaStream_t stream, AllGatherAlgo algo = AllGatherAlgo::kNone); + +/** + * @brief All-Reduce + * + * Reduces data arrays of length count in senddata using op operation, and + * leaves identical copies of result on each recvdata. + * + * In-place operation will happen if senddata == recvdata. + * + * @param comm: Communicator + * @param senddata: send data + * @param recvdata: recv data + * @param count: the number of elements, it is not nbytes. + * @param dtype: data type + * @param op：Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg + * @param stream: CUDA Stream + * @param algo: Algorithm + * @throw CommError: Throw CommError when an error is encountered. + */ +void allReduce(Comm_t comm, const void *senddata, void *recvdata, size_t count, ncclDataType_t dtype, + ncclRedOp_t op, cudaStream_t stream, AllReduceAlgo algo = AllReduceAlgo::kNone); + +/** + * @brief Whether is supported non-contiguous tensors + * @param comm: Communicator + * @param dtype: data type + * @param shape: tensor shape + * @param ndim: the ndim of tensor + * @param numel: the number of tensor + * @param op: Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg + * @return: supported + */ +bool allReduceStrideSupported(Comm_t comm, ncclDataType_t dtype, const int64_t *shape, int ndim, uint64_t numel, ncclRedOp_t op); + +/** + * @brief All-Reduce for non-contiguous tensors + * @param comm: Communicator + * @param senddata: send tensor + * @param recvdata: recv tensor + * @param stream: CUDA Stream + */ +void allReduceStride(Comm_t comm, const TensorDesc &senddata, TensorDesc &recvdata, cudaStream_t stream); + +/** + * @brief Send data from senddata to rank peer. + * + * Rank peer needs to call ncclRecv with the same datatype and the same count from this + * rank. This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations + * need to progress concurrently to complete. + * + * @param comm: Communicator + * @param senddata: send data + * @param count: the number of send elements, it is not nbytes. + * @param dtype: data type + * @param peer: the destination rank + * @param stream: CUDA Stream + * @param algo: Algorithm + * @throw CommError: Throw CommError when an error is encountered. + */ +void send(Comm_t comm, const void *senddata, size_t count, ncclDataType_t dtype, int peer, cudaStream_t stream, + SendAlgo algo = SendAlgo::kNone); + +/** + * @brief Receive data from rank peer into recvdata. + * + * Rank peer needs to call ncclSend with the same datatype and the same count to this + * rank. This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations + * need to progress concurrently to complete. + * + * @param comm: Communicator + * @param recvdata: recv data + * @param count: the number of recv elements, it is not nbytes. + * @param dtype: data type + * @param peer: source rank + * @param stream: CUDA Stream + * @param algo: Algorithm + * @throw CommError: Throw CommError when an error is encountered. + */ +void recv(Comm_t comm, void *recvdata, size_t count, ncclDataType_t dtype, int peer, cudaStream_t stream, + RecvAlgo algo = RecvAlgo::kNone); + +/** + * @brief Reduces data arrays of length count in senddata into recvdata using op operation. + * + * Recvdata may be NULL on all calls except for root device. + * root is the rank (not the CUDA device) where data will reside after the + * operation is complete. + * + * In-place operation will happen if senddata == recvdata. + * + * @param comm: Communicator + * @param senddata: send data + * @param recvdata: recv data + * @param count: the number of elements, it is not nbytes. + * @param dtype: data type + * @param op: Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg + * @param root: root rank + * @param stream: CUDA Stream + * @param algo: Algorithm + * @throw CommError: Throw CommError when an error is encountered. + */ +void reduce(Comm_t comm, const void *senddata, void *recvdata, size_t count, ncclDataType_t dtype, ncclRedOp_t op, + int root, cudaStream_t stream, ReduceAlgo algo = ReduceAlgo::kNone); + +/** + * @brief Broadcast + * + * Copies count values from root to all other devices. + * root is the rank (not the CUDA device) where data resides before the + * operation is started. + * + * In-place operation will happen if senddata == recvdata. + * + * @param comm: Communicator + * @param senddata: send data + * @param recvdata: recv data + * @param count: the number of elements, it is not nbytes. + * @param dtype: data type + * @param root: root rank + * @param stream: CUDA Stream + * @param algo: Algorithm + * @throw CommError: Throw CommError when an error is encountered. + */ +void broadcast(Comm_t comm, const void *senddata, void *recvdata, size_t count, ncclDataType_t dtype, int root, + cudaStream_t stream, BroadcastAlgo algo = BroadcastAlgo::kNone); + +/** + * + * @brief Reduce-Scatter + * + * Reduces data in senddata using op operation and leaves reduced result + * scattered over the devices so that recvdata on rank i will contain the i-th + * block of the result. + * Assumes sendcount is equal to nranks*recvcount, which means that senddata + * should have a size of at least nranks*recvcount elements. + * + * In-place operations will happen if recvdata == senddata + rank * recvcount. + * + * @param comm: Communicator + * @param senddata: send data + * @param recvdata: recv data + * @param recvcount: the number of recv elements, it is not nbytes. + * @param dtype: data type + * @param op：Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg + * @param stream: CUDA Stream + * @param algo: Algorithm + * @throw CommError: Throw CommError when an error is encountered. + */ +void reduceScatter(Comm_t comm, const void *senddata, void *recvdata, + size_t recvcount, ncclDataType_t dtype, ncclRedOp_t op, cudaStream_t stream, + ReduceScatterAlgo algo = ReduceScatterAlgo::kNone); + +/** + * @brief Send data from src_rank to dst_rank on src_rank process, recv data on dst_rank. + * + * @param comm: Communicator + * @param data: send data to dst rank if current rank is src_rank, recv data if current rank is dst_rank. + * @param count: the number of send/recv elements, it is not nbytes. + * @param dtype: data type + * @param src_rank: src rank + * @param dst_rank: dst rank + * @param stream: CUDA Stream + * @param algo: Algorithm + * @throw CommError: Throw CommError when an error is encountered. + */ +void p2p(Comm_t comm, void *data, size_t count, ncclDataType_t dtype, int src_rank, int dst_rank, cudaStream_t stream, + SendAlgo algo = SendAlgo::kNone); + +/** + * @brief The all ranks of communicator send senddata to dst_rank. + * + * @param comm: Communicator + * @param senddata: send data + * @param recvdatas: recv datas，it is two-dim array, shape: [WorldSize, RecvDataPointer], + * the first dim is host pointer，the second dim is GPU pointer, + * it can be nullptr when current rank is not dst rank. + * @param sendcount: the number of send elements, it is not nbytes. + * @param dst_rank: dst rank + * @param dtype: data type + * @param stream: CUDA Stream + * @param algo: Algorithm + * @throw CommError: Throw CommError when an error is encountered. + */ +void gather(Comm_t comm, const void *senddata, void **recvdatas, size_t sendcount, int dst_rank, ncclDataType_t dtype, + cudaStream_t stream, GatherAlgo algo = GatherAlgo::kNone); + + +}// namespace ixformer::comm diff --git a/ixformer_sdk/csrc/include/ixformer/comm/core/error.h b/ixformer_sdk/csrc/include/ixformer/comm/core/error.h new file mode 100644 index 00000000..c41276d8 --- /dev/null +++ b/ixformer_sdk/csrc/include/ixformer/comm/core/error.h @@ -0,0 +1,21 @@ +#pragma once + +#include +#include "status.h" + +namespace ixformer::comm { + +class CommError : public std::runtime_error { +public: + template + CommError(CommStatus error, const ERROR_STR str) : error_{error}, std::runtime_error(str) {} + + CommStatus status() { + return error_; + } + +private: + CommStatus error_; +}; + +}// namespace ixformer::comm diff --git a/ixformer_sdk/csrc/include/ixformer/comm/core/op_algo.h b/ixformer_sdk/csrc/include/ixformer/comm/core/op_algo.h new file mode 100644 index 00000000..491fc206 --- /dev/null +++ b/ixformer_sdk/csrc/include/ixformer/comm/core/op_algo.h @@ -0,0 +1,80 @@ +#pragma once + +#include + +namespace ixformer::comm { + +enum class AllGatherAlgo { + kNone, + kAuto, + kNCCL, + kNumAlgo +}; + +enum class AllReduceAlgo { + kNone, // None + kAuto, // 自动选择算法 + kAllGatherSum, // 针对小数据量的算法 + kBroadcastSum, // 针对小数据量的算法 + kRing, // Ring AllReduce + kQuant, // 对通讯算法进行量化，默认使用 kQuantL1 + kQuantL1, // 对通讯算法进行量化，优先使用量化算法以及最大保留精度，在部分 Size 性能不佳时，退化为 Auto 算法 + kQuantL2, // 对通讯算法进行量化，优先使用量化算法以及最大化速度，在部分 Size 性能不佳时，退化为 Auto 算法 + kQuantL1AllSize,// 对所有的 Size 都使用量化算法 + kQuantL2AllSize,// 对所有的 Size 都使用量化算法 + kNCCL, // 使用 NCCL + kStride, // 输入或输出的 Tensor 不是连续的 + kNumAlgo +}; + + +enum class BroadcastAlgo { + kNone, + kAuto, + kNCCL, + kNumAlgo +}; + +enum class GatherAlgo { + kNone, + kAuto, + kNCCL, + kNumAlgo +}; + +enum class SendAlgo { + kNone, + kAuto, + kNCCL, + kNumAlgo +}; + +typedef SendAlgo RecvAlgo; + +enum class ReduceAlgo { + kNone, + kAuto, + kNCCL, + kNumAlgo +}; + +enum class ReduceScatterAlgo { + kNone, + kAuto, + kNCCL, + kNumAlgo +}; + + +std::string to_string(AllGatherAlgo algo); +std::string to_string(AllReduceAlgo algo); +std::string to_string(BroadcastAlgo algo); +std::string to_string(GatherAlgo algo); +std::string to_string(SendAlgo algo); +std::string to_string(ReduceAlgo algo); +std::string to_string(ReduceScatterAlgo algo); + +template +Algo get_algo_from_str(const std::string &name); + +}// namespace ixformer::comm diff --git a/ixformer_sdk/csrc/include/ixformer/comm/core/status.h b/ixformer_sdk/csrc/include/ixformer/comm/core/status.h new file mode 100644 index 00000000..3c7d117b --- /dev/null +++ b/ixformer_sdk/csrc/include/ixformer/comm/core/status.h @@ -0,0 +1,22 @@ +#pragma once +#include "comm/core/common.h" + + +namespace ixformer::comm { + +enum CommStatus { + commSuccess, + commFail, + commCudaError, + commNcclError, + commInvalidArgument, + commUnsupported, + commInternalError, + commInvalidComm// maybe comm is nullptr +}; + + +std::string to_string(CommStatus status); + + +}// namespace ixformer::comm diff --git a/ixformer_sdk/csrc/include/ixformer/kernels/error.h b/ixformer_sdk/csrc/include/ixformer/kernels/error.h new file mode 100644 index 00000000..31060b1e --- /dev/null +++ b/ixformer_sdk/csrc/include/ixformer/kernels/error.h @@ -0,0 +1,22 @@ +#pragma once + +#include + +#include "status.h" + +namespace ixformer::kernels { + +class KernelError : public std::runtime_error { +public: + template + KernelError(KernelStatus error, const ERROR_STR str) : error_{error}, std::runtime_error(str) {} + + KernelStatus status() { + return error_; + } + +private: + KernelStatus error_; +}; + +}// namespace ixformer::kernels diff --git a/ixformer_sdk/csrc/include/ixformer/kernels/kernels.h b/ixformer_sdk/csrc/include/ixformer/kernels/kernels.h new file mode 100644 index 00000000..00c12486 --- /dev/null +++ b/ixformer_sdk/csrc/include/ixformer/kernels/kernels.h @@ -0,0 +1,2520 @@ +#pragma once + +#include "error.h" +#include "status.h" +#include "tensor.h" +#include +#include +#include +#include + + +namespace ixformer::kernels::infer { + + +/// ======================================================== +// Paged attention +// ======================================================== + +typedef enum { + KV_CACHE_FORMAT_STD, + KV_CACHE_FORMAT_NHD, + KV_CACHE_FORMAT_HND +} kvCacheFormat; + +/** + * @brief + * + * @tparam DType: input data type: half or bfloat16 + * @tparam IndexDType: index data type, int32 + * @tparam Format: KV_CACHE_FORMAT_STD or KV_CACHE_FORMAT_NHD or KV_CACHE_FORMAT_HND + * @param key: key, shape: [num_tokens,num_heads, head_size] + * @param value: value, shape: [num_tokens, num_heads, head_size] + * @param key_cache: key cache, half or bfloat16 + * 1. kv_cache_format == "STD", shape: [num_blocks, num_kv_heads, head_size // x, block_size, x] + * 2. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 3. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param value_cache: value cache, half or bfloat16 + * 1. kv_cache_format == "STD", shape: [num_blocks, num_kv_heads, head_size, block_size] + * 2. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 3. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param slot_mapping: The mapping position of the token in blocks. + * @param key_stride: key.stride(0) + * @param value_stride: value.stride(0) + * @param key_cache_stride: key_cache.stride(0) + * @param value_cache_stride: value_cache.stride(0) + * @param num_tokens: the number of tokens + * @param num_heads: the number of heads + * @param head_size: head size + * @param block_size: tokens of each page + * @param x: usually x = 16 / sizeof(DType) + * @param stream: CUDA Stream + */ +template +void paged_attention_cache_appended_f16_kernel( + const DType *key, + const DType *value, + DType *key_cache, + DType *value_cache, + const IndexDType *slot_mapping, + unsigned key_stride, + unsigned value_stride, + unsigned key_cache_stride, + unsigned value_cache_stride, + unsigned num_tokens, + unsigned num_heads, + unsigned head_size, + unsigned block_size, + unsigned x, + cudaStream_t stream); + +/** + * @brief + * + * @tparam DType: input data type: half or bfloat16 + * @tparam IndexDType: index data type, int32 + * @tparam Format: only support KV_CACHE_FORMAT_NHD or KV_CACHE_FORMAT_HND + * @param key: key, shape: [num_tokens, num_heads, head_size] + * @param value: value, shape: [num_tokens, num_heads, head_size] + * @param key_cache: key cache, int8, + * 1. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 2. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param value_cache: value cache, int8, + * 1. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 2. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param key_cache_scales: key cache scales, shape: [num_blocks, block_size] + * @param value_cache_scales: value cache scales, shape: [num_blocks, block_size] + * @param slot_mapping: The mapping position of the token in blocks. + * @param key_stride: key.stride(0) + * @param value_stride: value.stride(0) + * @param key_cache_stride: key_cache.stride(0) + * @param value_cache_stride: value_cache.stride(0) + * @param num_tokens: the number of tokens + * @param num_heads: the number of heads + * @param head_size: head size + * @param block_size: tokens of each page + * @param x: usually x = 16 / sizeof(DType) + * @param stream: CUDA Stream + */ +template +void paged_attention_cache_appended_i8_kernel( + const DType *key, + const DType *value, + int8_t *key_cache, + int8_t *value_cache, + DType *key_cache_scales, + DType *value_cache_scales, + const IndexDType *slot_mapping, + unsigned key_stride, + unsigned value_stride, + unsigned key_cache_stride, + unsigned value_cache_stride, + unsigned num_tokens, + unsigned num_heads, + unsigned head_size, + unsigned block_size, + unsigned x, + cudaStream_t stream); + +/** + * @brief + * + * @tparam DType: input data type: half or bfloat16 + * @tparam Format: KV_CACHE_FORMAT_STD or KV_CACHE_FORMAT_NHD or KV_CACHE_FORMAT_HND + * @param query: query, shape: [num_tokens, num_heads, head_size] + * @param key_cache: key cache, half or bfloat16 + * 1. kv_cache_format == "STD", shape: [num_blocks, num_kv_heads, head_size // 8, block_size, 8] + * 2. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 3. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param value_cache: value cache, half or bfloat16 + * 1. kv_cache_format == "STD", shape: [num_blocks, num_kv_heads, head_size // 8, block_size, 8] + * 2. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 3. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param block_tables: bloack tables, is used to store block + * @param seq_lens: shape: [num_tokens] + * @param alibi_slopes: alibi slopes + * @param out: output tensor, shape: [num_tokens, num_heads, head_size] + * @param aux_md: aux_max_expsum, is used to store intermediate values + * @param aux_o: aux_output, is used to store intermediate values + * @param max_seq_len: max seq len in a batch + * @param num_kv_heads: the number of kv heads + * @param num_heads: the number of query heads + * @param num_seqs: the number of seqs, num_seqs = query.size(0) + * @param head_size: head size + * @param block_size: tokens of each page + * @param max_num_blocks_per_seq: (MAX_SEQ_LEN + block_size - 1) // block_size + * @param q_stride: query.stride(0) + * @param kv_block_stride: key_cache.stride(0) + * @param scale: attention scale value + * @param use_sqrt_alibi: whether to use sqrt alibi + * @param stream: CUDA Stream + */ +template +void paged_attention_f16_algo0_kernel( + const DType *query, + const DType *key_cache, + const DType *value_cache, + const int *block_tables, + const int *seq_lens, + const float *alibi_slopes, + DType *out, + float *aux_md, + float *aux_o, + unsigned max_seq_len, + unsigned num_kv_heads, + unsigned num_heads, + unsigned num_seqs, + unsigned head_size, + unsigned block_size, + unsigned max_num_blocks_per_seq, + unsigned q_stride, + unsigned kv_block_stride, + float scale, + bool use_sqrt_alibi, + cudaStream_t stream); + +/** + * @brief + * + * @tparam DType: input data type: half or bfloat16 + * @tparam Format: KV_CACHE_FORMAT_STD or KV_CACHE_FORMAT_NHD or KV_CACHE_FORMAT_HND + * @param query: query, shape: [num_tokens, num_heads, head_size] + * @param key_cache: key cache, half or bfloat16 + * 1. kv_cache_format == "STD", shape: [num_blocks, num_kv_heads, head_size // 8, block_size, 8] + * 2. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 3. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param value_cache: value cache, half or bfloat16 + * 1. kv_cache_format == "STD", shape: [num_blocks, num_kv_heads, head_size // 8, block_size, 8] + * 2. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 3. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param block_tables: bloack tables, is used to store block + * @param seq_lens: shape: [num_tokens] + * @param alibi_slopes: alibi slopes + * @param out: output tensor, shape: [num_tokens, num_heads, head_size] + * @param aux_md: aux_max_expsum, is used to store intermediate values + * @param aux_o: aux_output, is used to store intermediate values + * @param max_seq_len: max seq len in a batch + * @param num_kv_heads: the number of kv heads + * @param num_heads: the number of query heads + * @param num_seqs: the number of seqs, num_seqs = query.size(0) + * @param head_size: head size + * @param block_size: block size + * @param max_num_blocks_per_seq: (MAX_SEQ_LEN + block_size - 1) // block_size + * @param q_stride: query.stride(0) + * @param kv_block_stride: key_cache.stride(0) + * @param scale: attention scale value + * @param use_sqrt_alibi: whether to use sqrt alibi + * @param stream: CUDA Stream + */ +template +void paged_attention_f16_algo1_kernel( + const DType *query, + const DType *key_cache, + const DType *value_cache, + const int *block_tables, + const int *seq_lens, + const float *alibi_slopes, + DType *out, + float *aux_md, + float *aux_o, + unsigned max_seq_len, + unsigned num_kv_heads, + unsigned num_heads, + unsigned num_seqs, + unsigned head_size, + unsigned block_size, + unsigned max_num_blocks_per_seq, + unsigned q_stride, + unsigned kv_block_stride, + float scale, + bool use_sqrt_alibi, + cudaStream_t stream); + +/** + * @brief + * + * @tparam DType: input data type: half or bfloat16 + * @tparam Format: KV_CACHE_FORMAT_NHD or KV_CACHE_FORMAT_HND + * @param query: query, shape: [num_tokens, num_heads, head_size] + * @param key_cache: key cache, int8, + * 1. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 2. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param value_cache: value cache, int8, + * 1. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 2. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param key_cache_scales: key cache scales, shape: [num_blocks, block_size] + * @param value_cache_scales: value cache scales, shape: [num_blocks, block_size] + * @param block_tables: bloack tables, is used to store block, shape:[num_seqs, max_num_blocks_per_seq] + * @param seq_lens: shape: [num_tokens] + * @param alibi_slopes: alibi slopes, shape:[num_heads] + * @param out: output tensor, shape: [num_tokens, num_heads, head_size] + * @param aux_md: aux_max_expsum, is used to store intermediate values + * @param aux_o: aux_output, is used to store intermediate values + * @param max_seq_len: max seq len in a batch + * @param num_kv_heads: the number of kv heads + * @param num_heads: the number of query heads + * @param num_seqs: the number of seqs, num_seqs = query.size(0) + * @param head_size: head size + * @param block_size: tokens of each page + * @param max_num_blocks_per_seq: (MAX_SEQ_LEN + block_size - 1) // block_size + * @param q_stride: query.stride(0) + * @param kv_block_stride: key_cache.stride(0) + * @param scale: attention scale value + * @param use_sqrt_alibi: whether to use sqrt alibi + * @param stream: CUDA Stream + */ +template +void paged_attention_i8_algo1_kernel( + const DType *query, + const int8_t *key_cache, + const int8_t *value_cache, + const DType *key_cache_scales, + const DType *value_cache_scales, + const int *block_tables, + const int *seq_lens, + const float *alibi_slopes, + DType *out, + float *aux_md, + float *aux_o, + unsigned max_seq_len, + unsigned num_kv_heads, + unsigned num_heads, + unsigned num_seqs, + unsigned head_size, + unsigned block_size, + unsigned max_num_blocks_per_seq, + unsigned q_stride, + unsigned kv_block_stride, + float scale, + bool use_sqrt_alibi, + cudaStream_t stream); + +/** + * @brief + * + * @tparam DType: input data type: half or bfloat16 + * @tparam Format: only support KV_CACHE_FORMAT_NHD or KV_CACHE_FORMAT_HND + * @param query: query, shape: [num_tokens, num_heads, head_size] + * @param paged_k_data: key_cache, + * 1. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 2. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] +* @param paged_v_data: value_cache, + * 1. kv_cache_format == "NHD", shape: [num_blocks, block_size, num_kv_heads, head_size] + * 2. kv_cache_format == "HND", shape: [num_blocks, num_kv_heads, block_size, head_size] + * @param paged_kv_indptr: the number of kv blocks in per token. shape: [num_tokens] + * @param paged_kv_indices: kv block index. shape: [num_blocks] + * @param paged_kv_last_page_len: shape: [num_tokens] + * @param alibi_slopes: alibi slopes, shape:[num_heads] + * @param out: output tensor, shape: [num_tokens, num_heads, head_size] + * @param aux_md: aux_max_expsum, is used to store intermediate values + * @param aux_o: aux_output, is used to store intermediate values + * @param max_seq_len: max seq len + * @param num_kv_heads: the number of kv heads + * @param num_qo_heads: the number of query heads + * @param num_seqs: the number of seqs + * @param head_size: head size + * @param page_size: block size + * @param q_stride: query.stride(0) + * @param kv_block_stride: kv_block.stride(0) + * @param scale: attention scale value + * @param use_sqrt_alibi: whether to use sqrt alibi + * @param stream: CUDA Stream + */ +template +void paged_attention_flashinfer_f16_kernel(const DType *query, + const DType *paged_k_data, + const DType *paged_v_data, + const int32_t *paged_kv_indptr, + const int32_t *paged_kv_indices, + const int32_t *paged_kv_last_page_len, + const float *alibi_slopes, + DType *out, + float *aux_md, + float *aux_o, + int32_t max_seq_len, + unsigned num_kv_heads, + unsigned num_qo_heads, + unsigned num_seqs, + unsigned head_size, + unsigned page_size, + unsigned q_stride, + unsigned kv_block_stride, + float scale, + bool use_sqrt_alibi, + cudaStream_t stream); + + + + +// ======================================================== +// MOE +// ======================================================== + +/** + * @brief + * + * @tparam T1: input type, half or bfloat16 + * @tparam T2: output type, half or bfloat16 + * @param A: The input tensor representing tokens with shape (num_tokens, K), + * where K is the feature dimension of each token. shape: [num_tokens, K] + * @param align_A: The input tensor representing tokens post padding with shape (pad_m, K), + * where pad_m is the total number of tokens post padding and K is the feature dimension of each token. + * @param B: The stacked MOE weight tensor with shape (E, N, K), + * where E is the number of experts, K is the input feature dimension, and N is the output feature dimension. + * @param C: The output cache tensor with shape (M, topk, N), where M is the total number of tokens post padding, + * topk is the number of times each token is repeated, and N is the output feature dimension. + * @param topk_weight: topk weight, shape: [num_tokens, topk] + * @param topk_ids: topk index, shape: [num_tokens, topk] + * @param sorted_token_ids: The tensor containing the sorted indices of tokens, + * repeated topk times and arranged by the expert index they are assigned to. + * shape: [topk_ids.numel() + num_experts * (block_size - 1)] + * @param expert_ids:The tensor containing the indices of the expert for each block. + * It determines which expert matrix from B should be used for each block in A. + * shape: [topk_ids.numel() + num_experts] + * @param m: the total number of tokens + * @param pad_m: the total number of tokens post padding + * @param n: B.size(1), the output feature dimension + * @param k: B.size(2), the feature dimension of each token + * @param top_k: topk + * @param block_size_m: BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix multiplication + * across different blocks processed by the same expert. + * @param mul_routed_weight: Whether to apply route weight + * @param stream: CUDA Stream + */ +template +void fused_moe(const T1 *A, T1 *align_A, const T1 *B, T1 *C, + const float *topk_weight, const int32_t *topk_ids, const int32_t *sorted_token_ids, + const int32_t *expert_ids, unsigned m, unsigned pad_m, unsigned n, unsigned k, + unsigned top_k, unsigned block_size_m, + bool mul_routed_weight, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, int8 + * @param A: The input tensor representing tokens with shape (num_tokens, K), + * where K is the feature dimension of each token + * @param align_A: The input tensor representing tokens post padding with shape (pad_m, K), + * where pad_m is the total number of tokens post padding and K is the feature dimension of each token. + * @param B: The stacked MOE weight tensor with shape (E, N, K), where E is the number of experts, + * K is the input feature dimension, and N is the output feature dimension. shape: [E, N, K] + * @param C: The output cache tensor with shape (M, topk, N), where M is the total number of tokens post padding, + * topk is the number of times each token is repeated, and N is the output feature dimension. + * @param topk_weight: topk weight, shape: [num_tokens, topk] + * @param topk_ids: topk index, shape: [num_tokens, topk] + * @param sorted_token_ids: The tensor containing the sorted indices of tokens, + * repeated topk times and arranged by the expert index they are assigned to. + * shape: [topk_ids.numel() + num_experts * (block_size - 1)] + * @param expert_ids: The tensor containing the indices of the expert for each block. + * It determines which expert matrix from B should be used for each block in A. + * shape: [topk_ids.numel() + num_experts] + * @param w_scale: B scale + * @param a_scale: A scale + * @param persistent: persistent + * @param expert_num: the number of experts + * @param m: the total number of tokens + * @param pad_m: the total number of tokens post padding + * @param n: B.size(1), the output feature dimension + * @param k: B.size(2), the feature dimension of each token + * @param top_k: topk + * @param block_size_m: BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix + * multiplication across different blocks processed by the same expert. + * @param mul_routed_weight: Whether to apply route weight + * @param input_extend: Determine whether the input tensor needs to be extended + * @param stream: CUDA Stream + * @param cuinfer_handle: CUINFER HANDLE + */ +template +void fused_moe_ixinfer(const int8_t *A, int8_t *align_A, const int8_t *B, T *C, + const float *topk_weight, const int32_t *topk_ids, + const int32_t *sorted_token_ids, const int32_t *expert_ids, + const float *w_scale, const float *a_scale, int64_t persistent, unsigned expert_num, + unsigned m, unsigned pad_m, unsigned n, unsigned k, unsigned top_k, unsigned block_size_m, + bool mul_routed_weight, bool input_extend, cudaStream_t stream, cuinferHandle_t cuinfer_handle); + +/** + * @brief + * + * @tparam T: input type, float + * @param gating_output: input tensor, shape: [num_tokens, num_experts] + * @param topk_weights: topk weights, shape: [num_tokens, topk] + * @param topk_indices: topk indices, shape: [num_tokens, topk] + * @param token_expert_indices: expert indices, shape: [num_tokens, topk] + * @param softmax_workspace: softmax workspace + * @param num_tokens: the number of tokens + * @param num_experts: the number of experts + * @param topk: topk + * @param renormalize: whether renormalize the result + * @param stream: CUDA Stream + */ +template +void moe_topk_softmax( + const T *gating_output, + T *topk_weights, + int *topk_indices, + int *token_expert_indices, + T *softmax_workspace, + int num_tokens, + int num_experts, + int topk, + bool renormalize, + cudaStream_t stream); + +/** + * @brief + * + * @tparam IN_DTYPE: gating_output type, half or bfloat16 + * @tparam INDEX_DTYPE: topk_indices type, int32 or int64 + * @param topk_weights: topk weights, shape: [num_tokens, topk] + * @param topk_indices: topk indices, shape: [num_tokens, topk] + * @param gating_output: input tensor, shape: [num_tokens, num_experts] + * @param bias: bias tensor for grouped topk, shape: [num_experts] + * @param num_tokens: the number of tokens + * @param num_experts: the number of experts + * @param topk: topk + * @param num_expert_group: num_expert_group + * @param topk_group: topk_group + * @param renormalize: whether renormalize the result + * @param scoring_func: scoring function for grouped topk + * @param stream: CUDA Stream + */ +template +void moe_grouped_topk( + float *topk_weights, + INDEX_DTYPE *topk_indices, + const IN_DTYPE *gating_output, + const IN_DTYPE *bias, + int num_tokens, int num_experts, int topk, + int num_expert_group, int topk_group, bool renormalize, + std::string scoring_func, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, int32_t + * @param topk_ids: topk index, shape: [num_tokens, topk] + * @param sorted_token_ids: The tensor containing the sorted indices of tokens, + * repeated topk times and arranged by the expert index they are assigned to. + * shape: [topk_ids.numel() + num_experts * (block_size - 1)] + * @param expert_ids: The tensor containing the indices of the expert for each block. + * It determines which expert matrix from B should be used for each block in A. + * shape: [topk_ids.numel() + num_experts] + * @param total_tokens_post_pad: the number of tokens + * @param aux_tokens_cnts: used for large num_experts + * @param aux_cumsum: used for large num_experts + * @param num_experts: the number of experts + * @param block_size: tokens of each page + * @param numel: topk_ids.numel() + * @param stream: CUDA Stream + */ +template +void moe_align_block_size(const T *topk_ids, + int32_t *sorted_token_ids, + int32_t *expert_ids, + int32_t *total_tokens_post_pad, + int32_t *aux_tokens_cnts, + int32_t *aux_cumsum, + int32_t num_experts, + int32_t block_size, + size_t numel, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param input: input tensor, shape: [outer_size, reduce_size, inner_size] + * @param mul_weights: broadcast mul before reduce sum, tensor shape: [outer_size, reduce_size] + * @param mask: control the validity of each vector, tensor shape: [outer_size, reduce_size] + * @param extra_residual: add on the final output, tensor shape: [outer_size, inner_size] + * @param out: output tensor, shape: [outer_size, inner_size] + * @param outer_size: outer_size + * @param reduce_size: reduce_size + * @param inner_size: inner_size + * @param in_stride: input.stride(1) + * @param out_stride: out.stride(0) + * @param scaling_factor: scaling factor for the output before residual + * @param stream: CUDA Stream + */ +template +void moe_output_reduce_sum( + const T *input, + const float *mul_weights, + const bool *mask, + const T *extra_residual, + T *out, + unsigned outer_size, + unsigned reduce_size, + unsigned inner_size, + unsigned in_stride, + unsigned out_stride, + float scaling_factor, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @tparam SCALE_T: smooth scales type, float32 or same as input + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param topk_ids: expert id for each tokens, shape: [num_tokens, topk] + * @param smooth_scales: smooth quant scales tensor for each experts, shape: [num_experts, hidden_size] + * @param dst_to_src: index of dst to src, shape: [num_tokens * topk] + * @param src_to_dst: index of src to dst, e.g. src_tensor[i] = dst_tensor[src_to_dst[i]]. shape: [num_tokens * topk] + * @param i8_outputs: output tensor shape: [dst_tokens, hidden_size] + * @param output_scales: scales tensor for output, shape: [dst_tokens] + * @param num_tokens: number tokens of input + * @param dst_tokens: the number of tokens after expansion + * @param hidden_size: hidden_size + * @param topk: topk for moe + * @param output_format: setting output format + * @param stream: CUDA Stream + */ +template +void moe_expand_input_dynamic_scaled_int8(const T *input, const int32_t *topk_ids, const SCALE_T *smooth_scales, + const int32_t *dst_to_src, const int32_t *src_to_dst, + int8_t *i8_outputs, float *output_scales, int num_tokens, + int dst_tokens, int hidden_size, int topk, int output_format, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input and output type, half or bfloat16 + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param dst_to_src: index of dst to src, shape: [num_tokens * topk] + * @param src_to_dst: index of src to dst, e.g. src_tensor[i] = dst_tensor[src_to_dst[i]]. shape: [num_tokens * topk] + * @param output: output tensor shape: [dst_tokens, hidden_size] + * @param num_tokens: number tokens of input + * @param dst_tokens: the number of tokens after expansion + * @param hidden_size: hidden_size + * @param topk: topk for moe + * @param stream: CUDA Stream + */ +template +void moe_expand_input(const T *input, const int32_t *dst_to_src, const int32_t *src_to_dst, T *output, + int num_tokens, int dst_tokens, int hidden_size, int topk, cudaStream_t stream); + +/** + * @brief + * + * @param topk_ids: expert id for each tokens, shape: [num_tokens, topk] + * @param src_dst: index of src to dst, e.g. src_tensor[i] = dst_tensor[src_to_dst[i]]. shape: [num_tokens * topk] + * @param dst_src: index of dst to src, shape: [num_tokens * topk] + * @param expert_sizes: the number of tokens allocated to each expert, shape: [num_experts] + * @param expand_tokens: the sum of expert_sizes, shape: [1] + * @param aux_tokens_cnts: used for large num_experts + * @param aux_cumsum: used for large num_experts + * @param num_experts: the numbers of num_experts overall + * @param start_expert_id: start expert id of the vaild expert interval + * @param end_expert_id: end expert id of the vaild expert interval [start_expert_id, end_expert_id) + * @param numel: size of topk_ids, the numbers of tokens + * @param stream: CUDA Stream + */ +void moe_compute_token_index( + int32_t *topk_ids, + int32_t *src_dst, + int32_t *dst_src, + int32_t *expert_sizes, + int32_t *expand_tokens, + int32_t *aux_tokens_cnts, + int32_t *aux_cumsum, + int32_t num_experts, + int32_t start_expert_id, + int32_t end_expert_id, + size_t numel, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @tparam SCALE_T: smooth scales type, float32 or same as input + * @tparam BIAS_T: bias type, float32 or same as input + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param bias: bias tensor, shape: [num_experts, hidden_size] + * @param smooth_scales: smooth quant scales tensor for each experts, shape: [num_experts, hidden_size // 2] if act_type==swiglu else [num_experts, hidden_size] + * @param dst_to_src: index of dst to src, shape: [num_tokens * topk] + * @param topk_ids: expert id for each tokens, shape: [num_tokens] + * @param out: output tensor, shape: [num_tokens, hidden_size // 2] if act_type==swiglu else [num_tokens, hidden_size] + * @param output_scales: scales tensor for output, shape: [num_tokens] + * @param act_type: str activation type. Options include gelu, silu, and swiglu. + * @param num_tokens: number tokens of input + * @param hidden_size: hidden_size + * @param output_format: setting output format + * @param stream: CUDA Stream + */ +template +void activation_dynamic_scaled_int8( + const T *input, const BIAS_T *bias, + const SCALE_T *smooth_scales, const int32_t *dst_to_src, + const int32_t *topk_ids, int8_t *out, float *output_scales, + std::string act_type, unsigned num_tokens, unsigned hidden_size, int output_format, cudaStream_t stream); + + + + +// ======================================================== +// Dynamic INT8 +// ======================================================== +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @tparam ST: smooth_scales type, same as T or be float32 + * @param input: input tensor, shape: [num_token, hidden_size] + * @param smooth_scales: input smooth scale tensor, shape: [hidden_size] + * @param out: output tensor, shape: [num_token, hidden_size] + * @param scale_output: output scale tensor, shape: [num_token] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param stream: CUDA Stream + */ +template +void dynamic_scaled_quant_smoothquant(const T *input, const ST *smooth_scales, int8_t *out, float *scale_output, + int num_tokens, int hidden_size, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param smooth_scales: input smooth scale tensor, shape: [hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output: output scale tensor, shape: [num_tokens] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param stream: CUDA Stream + */ +template +void silu_and_mul_smoothquant(const T *input, const T *smooth_scales, int8_t *out, float *scale_output, + int num_tokens, const int hidden_size, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param weight: weight, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param smooth_scales: input smooth scale tensor, shape: [hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output: output scale tensor, shape: [num_tokens] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void rmsnorm_smoothquant(const T *input, const T *weight, + const T *fused_bias, const ST *smooth_scales, + int8_t *out, float *scale_output, + int num_tokens, int hidden_size, int in_stride, + float eps, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param weight: weight, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output: output scale tensor, shape: [num_tokens] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void rmsnorm_dynamic_int8(const T *input, const T *weight, const T *fused_bias, + int8_t *out, float *scale_output, + int num_tokens, int hidden_size, int in_stride, + float eps, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @tparam ST: smooth_scales type, same as T or be float32 + * @tparam IS_POST: norm type, post or pre + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param residual: residual tensor, shape: [num_tokens, hidden_size] + * @param weight: weight, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param smooth_scales: input smooth scale tensor, shape: [hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output: output scale tensor, shape: [num_tokens] + * @param residual_output: residual_output[optional], shape: [num_tokens, hidden_size] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param resi_stride: int, the stride of dim "HiddenSize" to support non-contiguous residual + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void residual_rmsnorm_smoothquant(const T *input, T *residual, const T *weight, + const T *fused_bias, const ST *smooth_scales, + int8_t *out, float *scale_output, T *residual_output, + int num_tokens, int hidden_size, int in_stride, int resi_stride, + float eps, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @tparam IS_POST: norm type, post or pre + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param residual: residual tensor, shape: [num_tokens, hidden_size] + * @param weight: weight, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output: output scale tensor, shape: [num_tokens] + * @param residual_output: residual_output[optional], shape: [num_tokens, hidden_size] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param resi_stride: int, the stride of dim "HiddenSize" to support non-contiguous residual + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void residual_rmsnorm_dynamic_int8(const T *input, T *residual, const T *weight, const T *fused_bias, + int8_t *out, float *scale_output, T *residual_output, + int num_tokens, int hidden_size, int in_stride, int resi_stride, + float eps, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @tparam ST: smooth_scales type, same as T or be float32 + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param scale: weight, shape: [hidden_size] + * @param bias: bias, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param smooth_scales: input smooth scale tensor, shape: [hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output: output scale tensor, shape: [num_tokens] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void layernorm_smoothquant(const T *input, const T *scale, const T *bias, + const T *fused_bias, const ST *smooth_scales, + int8_t *out, float *scale_output, + int num_tokens, int hidden_size, int in_stride, + float eps, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param scale: weight, shape: [hidden_size] + * @param bias: bias, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output: output scale tensor, shape: [num_tokens] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void layernorm_dynamic_int8(const T *input, const T *scale, const T *bias, const T *fused_bias, + int8_t *out, float *scale_output, + int num_tokens, int hidden_size, int in_stride, + float eps, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @tparam ST: smooth_scales type, same as T or be float32 + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param residual: residual tensor, shape: [num_tokens, hidden_size] + * @param scale: weight, shape: [hidden_size] + * @param bias: bias, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param smooth_scales: input smooth scale tensor, shape: [hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output: output scale tensor, shape: [num_tokens] + * @param residual_output: residual_output[optional], shape: [num_tokens, hidden_size] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param resi_stride: int, the stride of dim "HiddenSize" to support non-contiguous residual + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void residual_layernorm_smoothquant(const T *input, T *residual, + const T *scale, const T *bias, + const T *fused_bias, const ST *smooth_scales, + int8_t *out, float *scale_output, T *residual_output, + int num_tokens, int hidden_size, + int in_stride, int resi_stride, + float eps, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param residual: residual tensor, shape: [num_tokens, hidden_size] + * @param scale: weight, shape: [hidden_size] + * @param bias: bias, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output: output scale tensor, shape: [num_tokens] + * @param residual_output: residual_output[optional], shape: [num_tokens, hidden_size] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param resi_stride: int, the stride of dim "HiddenSize" to support non-contiguous residual + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void residual_layernorm_dynamic_int8(const T *input, T *residual, + const T *scale, const T *bias, const T *fused_bias, + int8_t *out, float *scale_output, T *residual_output, + int num_tokens, int hidden_size, + int in_stride, int resi_stride, + float eps, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param scale1: weight, shape: [hidden_size] + * @param bias1: bias, shape: [hidden_size] + * @param smooth_scales1: input smooth scale tensor, shape: [hidden_size] + * @param scale2: weight, shape: [hidden_size] + * @param bias2: bias, shape: [hidden_size] + * @param smooth_scales2: input smooth scale tensor, shape: [hidden_size] + * @param output1: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output1: output scale tensor, shape: [num_tokens] + * @param output2: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output2: output scale tensor, shape: [num_tokens] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void layernorm_2sb_smoothquant(const T *input, + const T *scale1, const T *bias1, const T *smooth_scales1, + const T *scale2, const T *bias2, const T *smooth_scales2, + int8_t *output1, float *scale_output1, + int8_t *output2, float *scale_output2, + int num_tokens, int hidden_size, + float eps, cudaStream_t stream); + +/** + * @brief + * + * @tparam T + * @param input: input tensor, shape: [num_tokens, hidden_size] + * @param residual: residual tensor, shape: [num_tokens, hidden_size] + * @param scale1: weight, shape: [hidden_size] + * @param bias1: bias, shape: [hidden_size] + * @param smooth_scales1: input smooth scale tensor, shape: [hidden_size] + * @param scale2: weight, shape: [hidden_size] + * @param bias2: bias, shape: [hidden_size] + * @param smooth_scales2: input smooth scale tensor, shape: [hidden_size] + * @param output1: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output1: output scale tensor, shape: [num_tokens] + * @param output2: output tensor, shape: [num_tokens, hidden_size] + * @param scale_output2: output scale tensor, shape: [num_tokens] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void layernorm_2sb_residual_smoothquant(const T *input, T *residual, + const T *scale1, const T *bias1, const T *smooth_scales1, + const T *scale2, const T *bias2, const T *smooth_scales2, + int8_t *output1, float *scale_output1, + int8_t *output2, float *scale_output2, + int num_tokens, int hidden_size, + float eps, cudaStream_t stream); + + + + +// ======================================================== +// Lightllm +// ======================================================== + +/** + * @brief + * + * @tparam T: input type, float + * @param logits: apply_penalty input, shape: [batch, vocab_size] + * @param presence_penalty: Penalty term that controls whether the word exists. shape: [batch] + * @param freqency_penalty: Used to control the overall frequency of words in the generated text. shape: [batch] + * @param p_token_ids: The id corresponding to per token in the vocabulary，shape: [num_tokens] + * @param p_token_counts: The counts corresponding to per token. shape: [num_tokens] + * @param p_cumsum_seq_len: The cumulative value of seq_len in a batch. shape: [batch+1] + * @param p_max_len_in_batch: The maximum length of seq in a batch + * @param batch: Batch Size + * @param vocab_size: vocabulary size + * @param stream: CUDA Stream + */ +template +void lightllm_apply_penalty(T *logits, const T *presence_penalty, const T *freqency_penalty, + const int *p_token_ids, const int *p_token_counts, + const int *p_cumsum_seq_len, int p_max_len_in_batch, + int batch, int vocab_size, cudaStream_t stream); + + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param key_cache: key cache, shape: [num_tokens, num_kv_heads, head_size] + * @param b_mem_idx: Index of the destination location corresponding to the token. shape: [num_tokens]. + * @param out: output tensor, shape: [max_tokens, num_kv_heads, head_size] + * @param num_tokens: the number of tokens. + * @param num_heads: num_kv_heads. + * @param headdim: head_size + * @param stream: CUDA Stream + */ +template +void lightllm_destindex_copy_kv( + const T *key_cache, + const int *b_mem_idx, + T *out, + int num_tokens, + int num_heads, + int headdim, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half + * @param input: input tensor, shape: [num_tokens, head_num, head_dim] + * @param cos: shape: [num_tokens, 1, head_dim // 2 // 2] + * @param sin: shape: [num_tokens, 1, head_dim //2 //2] + * @param num_tokens: the number of tokens. + * @param head_num: the number of head. + * @param head_dim: head_size + * @param rot_dim: rot_dim = cos.size(-1) + * @param stream: CUDA Stream + */ + +template +void lightllm_glm2_rope(T *input, const T *cos, const T *sin, int num_tokens, + int head_num, int head_dim, int rot_dim, cudaStream_t stream); +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param out: output tensor, shape: [batch, head_num, head_dim] + * @param partition_size: partition size + * @param exp_sums: shape: [batch, num_heads, max_num_partitions] + * @param max_logits: shape: [batch, num_heads, max_num_partitions] + * @param tmp_out: shape: [batch, num_heads, max_num_partitions,head_size] + * @param query: shape: [batch,head_num,head_dim] + * @param key_cache: key cache. shape: [max_num_tokens, head_num_kv, head_dim] + * @param value_cache: value cache. shape: [max_num_tokens, head_num_kv, head_dim] + * @param scale:The scaling of QK^T before applying softmax. + * @param reg_to_tokens: shape: [max_requset,max_tokens] + * @param b_req_idx: request index in a batch, shape: [batch] + * @param b_seq_len: seq len in a batch. shape: [batch] + * @param q_stride: query.stride(0) + * @param kv_token_stride: key_cache.stride(0) + * @param kv_head_stride: key_cache.stride(1) + * @param max_context_len_cur_batch: b_seq_len.max() + * @param num_heads: the number of query head. + * @param num_kv_head: the number of kv head. + * @param batch: batch size + * @param stream:: CUDA Stream + */ +template +void lightllm_token_attention( + T *out, int64_t partition_size, float *exp_sums, + float *max_logits, T *tmp_out, + const T *query, + const T *key_cache, + const T *value_cache, + float scale, + const int *reg_to_tokens, + const int *b_req_idx, + const int *b_seq_len, + int q_stride, + int kv_token_stride, + int kv_head_stride, + int max_context_len_cur_batch, int num_heads, int num_kv_head, + int batch, cudaStream_t stream); + + + + +// ======================================================== +// Quant +// ======================================================== + +typedef enum { + QUANT_AWQ, + QUANT_GPTQ, + QUANT_INT8, + QUANT_NF4, + QUANT_FP4 +} QuantType; + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param input: input tensor, shape: [m, k] + * @param i8_input: int8 quant output, shape: [m, k] + * @param input_scales: scale + * @param is_dynamic: whether to use dynamic scale + * @param input_channel: input row + * @param output_channel: i8_input col + * @param stream:CUDA Stream + */ +template +void scaled_int8_quant(const T *input, int8_t *i8_input, float *input_scales, bool is_dynamic, int input_channel, int output_channel, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T1: output type, half or bfloat16 + * @tparam T2: output type after pack 4B, half2 or bfloat162 + * @tparam TYPE: QUANT_NF4 or QUANT_FP4 + * @param qweights: quant weight, uint8, shape: [output_channel, input_channel // 2] + * @param scales: scale, float, shape: [output_channel * input_channel // g] + * @param out: output tensor, shape: [input_channel, output_channel] + * @param input_channel: row + * @param output_channel: column + * @param group_size: group size + * @param stream: CUDA Stream + */ +template +void weight_dequant_float4(const unsigned char *qweights, const float *scales, T1 *out, + unsigned input_channel, unsigned output_channel, unsigned group_size, cudaStream_t stream); +/** + * @brief + * + * @tparam T: output type, half or bfloat16 + * @param qweights: quant weight, int32, shape: [input_channel // (32 / bits), output_channel] + * @param scales: scale, half or bfloat16, shape: [input_channel // g, output_channel] + * @param zeros: quant zeros, int32, shape: [input_channel // g, output_channel // (32 / bits)] + * @param g_idx: g_idx + * @param out: dequant output tensor, shape: [input_channel, output_channel] + * @param input_channel: output row + * @param bits: weight bits + * @param output_channel: output col + * @param group_size: group size + * @param deq_mode: dequant mode, + * 0: don't use g_idx + * 1: exllama with g_idx (g_idx has been argsort) + * 2: general with g_idx (g_idx mapping group_index for each input channel) + * @param stream: CUDA Stream + */ +template +void weight_dequant_gptq(const int *qweights, const T *scales, const int *zeros, const int32_t *g_idx, T *out, + int bits, unsigned input_channel, unsigned output_channel, unsigned group_size, int deq_mode, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, float, half, bfloat16 + * @tparam TYPE: QUANT_INT8 or QUANT_NF4 or QUANT_FP4 + * @param code: quantiztion map + * @param A: input tensor, shape: [row, col] + * @param absmax: shape: [row] + * @param out: output tensor + * @param rand: only support "None" + * @param rand_offset: only support 0 + * @param blocksize: block size + * @param n: total size of A + */ +template +void quantize_block_wise(const float *code, const T *A, float *absmax, unsigned char *out, const float *rand, + int rand_offset, int blocksize, int n, cudaStream_t stream); + +/** + * @brief + * + * @tparam T1: input type, half or bfloat16 + * @tparam T2: Intermediate variable type, half2 or bfloat162 + * @tparam use_ex: whether to use_exllama + * @param input: input tensor, shape: [bs, input_channel] + * @param scales: scale value, half or bfloat16, shape: [input_channel // g, output_channel] + * @param qweights: quant weight, int32, shape: [input_channel // 8, output_channel] + * @param qzeros: quant zero, int32, shape: [input_channel // g, output_channel // 8] + * @param bias: shape: [output_channel] + * @param g_idx:int32, shape: [input_channel] + * @param out: output tensor, shape: [bs, output_channel] + * @param bs: input row + * @param input_channel: input col + * @param output_channel: output col + * @param group_size: group size + * @param bits: quant bits + * @param stream: CUDA Stream + */ +template +void quantized_linear_int4_gptq(const T1 *input, const T1 *scales, const unsigned *qweights, const unsigned *qzeros, + T1 *bias, const int32_t *g_idx, T1 *out, unsigned bs, unsigned input_channel, unsigned output_channel, + unsigned group_size, unsigned bits, cudaStream_t stream); + +/** + * @brief + * + * @tparam T1: input type, half or bfloat16 + * @tparam T2: Intermediate variable type, half2 or bfloat162 + * @param input: input tensor, shape: [bs, input_channel] + * @param scales: scale value, half or bfloat16, shape: [input_channel // g, output_channel] + * @param qweights: quant weight, int32, shape: [input_channel // (32/BITS), output_channel] + * @param qzeros: quant zero, int32, shape: [input_channel // g, output_channel // (32/BITS)] + * @param bias: shape: [output_channel] + * @param g_idx:int32, shape: [input_channel] + * @param out: output tensor, shape: [bs, output_channel] + * @param bs: input row + * @param input_channel: input col + * @param output_channel: output col + * @param group_size: group size + * @param use_ex: wheather use exllama + * @param stream: CUDA Stream + */ +template +void quantized_linear_int8_gptq(const T1 *input, const T1 *scales, const unsigned *qweights, const unsigned *qzeros, + const T1 *bias, const int *g_idx, T1 *out, unsigned bs, unsigned input_channel, unsigned output_channel, unsigned group_size, bool use_ex, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T1: input type, half or bfloat16 + * @tparam T2: Intermediate variable type, half2 or bfloat162 + * @tparam quant_type: fp4 or nf4 + * @param input: input tensor, shape: [bs, input_channel] + * @param scales: scale value, float, shape: [output_channel * input_channel // g] + * @param qweights: quant weight, unint8, shape: [output_channel * input_channel // 2, 1] + * @param bias: shape: [output_channel] + * @param out: output tensor, shape: [bs, output_channel] + * @param bs: input row + * @param input_channel: input col + * @param output_channel: output col + * @param group_size: group size + * @param stream: CUDA Stream + */ +template +void quantized_linear_float4(const T1 *input, const float *scales, const unsigned char *qweights, const T1 *bias, T1 *out, + unsigned bs, unsigned input_channel, unsigned output_channel, unsigned group_size, cudaStream_t stream); + + + + +// ======================================================== +// act and mul +// ======================================================== + +/** + * @brief gelu_and_mul + * + * @tparam T: input type, half or bfloat16 or float + * @param input: gelu_and_mul input, shape: [num_tokens, 2 * hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param gate_first: bool type, Deciding whether the gelu function should be applied to the first half + * or the second half of the input + * @param stream: CUDA Stream + */ +template +void gelu_and_mul(const T *input, T *out, + int num_tokens, int hidden_size, bool gate_first, cudaStream_t stream); + +/** + * @brief gelu_tanh_and_mul + * + * @tparam T: input type, half or bfloat16 or float + * @param input: gelu_tanh_and_mul input, shape: [num_tokens, 2 * hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param stream: CUDA Stream + */ +template +void gelu_tanh_and_mul(const T *input, T *out, + int num_tokens, int hidden_size, cudaStream_t stream); + +/** + * @brief silu_and_mul + * + * @tparam T: input type, half or bfloat16 or float + * @param input: silu_and_mul input, shape: [num_tokens, 2 * hidden_size] + * @param out: output tensor, shape: [num_tokens, hidden_size] + * @param num_tokens: the number of tokens + * @param hidden_size: HiddenSize + * @param stream: CUDA Stream + */ +template +void silu_and_mul(const T *input, T *out, + int num_tokens, int hidden_size, cudaStream_t stream); + + +// ======================================================== +// LayerNorm +// ======================================================== + +/** + * @brief layernorm + * + * @tparam T: input type, half or bfloat16 + * @param input: layernorm input, shape: [batch_count * seq_len, hidden_size] + * @param scale: weight, shape: [hidden_size] + * @param bias: bias, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param out: output tensor, shape: [batch_count * seq_len, hidden_size] + * @param batch_tokens: int, Batch * InputTokens + * @param hidden_size: int, HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void layernorm(const T *input, const T *scale, const T *bias, const T *fused_bias, + T *out, int batch_tokens, int hidden_size, int in_stride, + float eps, cudaStream_t stream); + +/** + * @brief layernorm_residual + * + * @tparam T: input type, half or bfloat16 + * @tparam IS_POST: bool type, post-layernorm(true) or pre-layernorm(false) + * @param input: layernorm input, shape: [batch_count * seq_len, hidden_size] + * @param residual: residual tensor, shape: [batch_count * seq_len, hidden_size] + * @param scale: weight, shape: [hidden_size] + * @param bias: bias, shape: [hidden_size] + * @param fused_bias: fused_bias[optional], shape: [hidden_size] + * @param output: output[optional], shape: [batch_count * seq_len, hidden_size] + * @param residual_output: residual_output[optional], shape: [batch_count * seq_len, hidden_size] + * @param alpha: float, residual scale factor + * @param batch_tokens: int, Batch * InputTokens + * @param hidden_size: int, HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param resi_stride: int, the stride of dim "HiddenSize" to support non-contiguous residual + * @param eps: float, a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void layernorm_residual(T *input, T *residual, + const T *scale, const T *bias, + const T *fused_bias, + T *output, T *residual_output, + float alpha, int batch_tokens, int hidden_size, + int in_stride, int resi_stride, + float eps, cudaStream_t stream); +/** + * @brief layernorm_2sb + * + * @tparam T: input type, half or bfloat16 + * @param input: layernorm input, shape: [batch_count * seq_len, hidden_size] + * @param scale1: the first weight, shape: [hidden_size] + * @param bias1: the first bias, shape: [hidden_size] + * @param scale2: the second weight, shape: [hidden_size] + * @param bias2: the second bias, shape: [hidden_size] + * @param eps: float, a value added to the denominator for numerical stability + * @param output1: the first output tensor, shape: [batch_count * seq_len, hidden_size] + * @param output2: the second output tensor, shape: [batch_count * seq_len, hidden_size] + * @param batch_tokens: int, Batch * InputTokens + * @param hidden_size: int, HiddenSize + * @param stream: CUDA Stream + */ +template +void layernorm_2sb(const T *input, + const T *scale1, const T *bias1, + const T *scale2, const T *bias2, + float eps, + T *output1, T *output2, + int batch_tokens, int hidden_size, cudaStream_t stream); + +/** + * @brief layernorm + residual + 2sb + * + * @tparam T: input type, half or bfloat16 + * @param input: layernorm input, shape: [batch_count * seq_len, hidden_size] + * @param residual: residual tensor, shape: [batch_count * seq_len, hidden_size] + * @param scale1: the first weight, shape: [hidden_size] + * @param bias1: the first bias, shape: [hidden_size] + * @param scale2: the second weight, shape: [hidden_size] + * @param bias2: the second bias, shape: [hidden_size] + * @param eps: float, a value added to the denominator for numerical stability + * @param output1: the first output tensor, shape: [batch_count * seq_len, hidden_size] + * @param output2: the second output tensor, shape: [batch_count * seq_len, hidden_size] + * @param batch_tokens: int, Batch * InputTokens + * @param hidden_size: int, HiddenSize + * @param stream: CUDA Stream + */ +template +void layernorm_residual_2sb(const T *input, T *residual, + const T *scale1, const T *bias1, + const T *scale2, const T *bias2, + float eps, + T *output1, T *output2, + int batch_tokens, int hidden_size, cudaStream_t stream); + + +// ======================================================== +// RMS Norm +// ======================================================== + +/** + * @brief RMS Norm + * @tparam T: input type, half or bfloat16 + * @param input: RMS Norm input, shape: [Batch, InputTokens, HiddenSize] or [Batch * InputTokens, HiddenSize] + * @param weight: RMS Norm weight tensor, shape: [HiddenSize] + * @param fused_bias: fused_bias tensor[optional], shape: [HiddenSize] + * @param out: output tensor, shape: [Batch, InputTokens, HiddenSize] or [Batch * InputTokens, HiddenSize] + * @param batch_tokens: Batch * InputTokens + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param eps: a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void rms_norm(const T *input, const T *weight, const T *fused_bias, T *out, + int batch_tokens, int hidden_size, int in_stride, + float eps, cudaStream_t stream); + +/** + * @brief RMS Norm + Residual + * @tparam T: input type, half or bfloat16 + * @tparam IS_POST: bool type, post-layernorm(true) or pre-layernorm(false) + * @param input: RMS Norm input, shape: [Batch, InputTokens, HiddenSize] or [Batch * InputTokens, HiddenSize] + * @param residual: residual tensor, shape: [Batch, InputTokens, HiddenSize] or [Batch * InputTokens, HiddenSize] + * @param weight: RMS Norm weight tensor, shape: [HiddenSize] + * @param fused_bias: fused_bias tensor[optional], shape: [HiddenSize] + * @param output: output tensor, shape: [Batch, InputTokens, HiddenSize] or [Batch * InputTokens, HiddenSize] + * @param residual_output: residual output tensor[optional], shape: [Batch, InputTokens, HiddenSize] or [Batch * InputTokens, HiddenSize] + * @param batch_tokens: Batch * InputTokens + * @param alpha: float, residual scale factor + * @param hidden_size: HiddenSize + * @param in_stride: int, the stride of dim "HiddenSize" to support non-contiguous input + * @param resi_stride: int, the stride of dim "HiddenSize" to support non-contiguous residual + * @param eps: a value added to the denominator for numerical stability + * @param stream: CUDA Stream + */ +template +void rms_norm_residual(T *input, T *residual, const T *weight, + const T *fused_bias, T *output, T *residual_output, + int batch_tokens, float alpha, int hidden_size, int in_stride, int resi_stride, + float eps, cudaStream_t stream); + +// ======================================================== +// Softmax +// ======================================================== + +/** + * @brief softmax_2d + * + * @tparam T: input type, half + * @param input: softmax_2D input, shape: [*], where * means, any number of additional dimensions + * @param out:output tensor, same shape as input + * @param outer_dim: the product of all dimensions of input except the last dim + * @param inner_dim: the value is input.size(input.dim()-1) + * @param stream: CUDA Stream + */ +template +void softmax_2d(const T *input, T *out, int outer_dim, int inner_dim, + cudaStream_t stream); + +/** + * @brief fast_softmax_forwardimp + * + * @tparam T: input type, half,shape: [*], where * means, any number of additional dimensions + * @param stream: CUDA Stream + * @param input: fast_softmax input + * @param out: output tensor, same shape as input + * @param outer_dim: The product of all dimensions of input except the last dim + * @param inner_dim: the value is input.size(input.dim()-1) + */ +template +void fast_softmax_forwardimp(const T *input, T *out, int outer_dim, int inner_dim, cudaStream_t stream); + +// ======================================================== +// Add +// ======================================================== + +/** + * @brief element wise add + * + * @tparam T input type, half or bfloat16 or float + * @param A: input tensor, shape: (...) + * @param B: other tensor, shape: (...) same as A + * @param C: out tensor, shape: (...) same as A + * @param m: default = 1 + * @param n: A.numel() + * @param stream: CUDA Stream + */ +template +void add(const T *A, const T *B, T *C, int m, int n, cudaStream_t stream); + +// ======================================================== +// GroupNorm +// ======================================================== + +/** + * @brief groupnorm_ixinfer + * @tparam T: input type, half + * @param input: groupnorm_ixinfer input, shape: [N, C, H, W] or [N, H, W, C] or [N, C, HW] where C = num_channels + * @param scale: weight, shape: [C] + * @param bias: bias, shape: [C] + * @param out: output tensor, shape: [N, C, H, W] or [N, H, W, C] or [N, C, HW] where C = num_channels + * @param batch: Batch Size = N + * @param hw: Product of H and W + * @param num_channel: the number of channel, the value is C + * @param num_group: number of groups to separate the channels into + * @param eps: a value added to the denominator for numerical stability + * @param is_nhwc: bool type. NHWC or NCHW + * @param act_type: 0 or 1, if act_type=1, use silu; if act_type=0, no activate + * @param stream: CUDA Stream + */ +template +void groupnorm_ixinfer(const T *input, const T *scale, const T *bias, T *out, int batch, int hw, + int num_channel, int num_group, float eps, bool is_nhwc, int act_type, cudaStream_t stream); + + + + +// ======================================================== +// TGI +// ======================================================== +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param logits: input tensor, shape: [num_tokens, vocab_size] + * @param index : the indices of elements to gather + * @param out: output tensor + * @param n: The number of elements in index tensor + * @param vocab_size: vocabulary size + * @param stream: CUDA Stream + */ +template +void tgi_gather_prefill_logprobs(const T *logits, const int32_t *index, T *out, + int n, int vocab_size, cudaStream_t stream); + +/** + * @brief + * + * @tparam T1: input type, half or bfloat16 or float + * @tparam IS_NEOX: Determine whether to use Neox, that is, whether to use interleaved + * @param query1: The first half of the query tensor in last dimension, shape: [num_tokens, num_heads, head_size //2] + * @param query2: The second half of the query tensor in last dimension, shape: [num_tokens, num_heads, head_size //2] + * @param cos: applied in query1, shape: [max_position, 1, head_size //2] + * @param sin: applied in query2, shape: [max_position, 1, head_size //2] + * @param out1: The first half of output tensor in last dimension, shape: [num_tokens, num_heads, head_size //2] + * @param out2: The second half of output tensor in last dimension, shape: [num_tokens, num_heads, head_size //2] + * @param rot_dim: cos.size(2) + * @param query1_stride: query1.stride(0) + * @param num_tokens: the number of tokens + * @param num_heads: the number of heads + * @param head_size: head size + * @param stream: CUDA Stream + */ +template +void tgi_rotary_embedding_neox(const T1 *query1, + const T1 *query2, + const T1 *cos, + const T1 *sin, + T1 *out1, + T1 *out2, + int rot_dim, int query1_stride, + int num_tokens, int num_heads, int head_size, bool is_neox, + cudaStream_t stream); + + + + +// ======================================================== +// VLLM +// ======================================================== +/** + * @brief + * + * @tparam T: input type, half or bfloat16 or float + * @param key_cache: key cache, shape:[[num_blocks, num_kv_heads, block_size, head_size],...] + * @param value_cache: value cache, shape:[[num_blocks, num_kv_heads, block_size, head_size],...] + * @param block_mapping: shape: [num_tokens, 2] + * @param num_layers: the number of layers in a model. + * @param num_pairs: The number of tokens to be mapped. num_pairs = block_mapping.size(0) + * @param numel_per_block: the number of elements in per block, num_kv_heads* block_size*head_size + * @param stream: CUDA Stream + */ + +template +void vllm_copy_blocks( + int64_t *key_cache, + int64_t *value_cache, + const int64_t *block_mapping, + int num_layers, int num_pairs, + int numel_per_block, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 or float + * @param key: key. shape: [num_tokens, num_heads, head_size] + * @param value: value. shape: [num_tokens, num_heads, head_size] + * @param key_cache: key cache, shape: [num_blocks, num_heads, head_size//8, block_size, 8] + * @param value_cache: value cache. shape: [num_blocks, num_heads, head_size//8, block_size, 8] + * @param slot_mapping: The mapping position of the token in blocks. shape: [num_tokens] + * @param key_stride: key.stride(0) + * @param value_stride: value.stride(0) + * @param num_heads: the number of heads + * @param head_size: head size + * @param block_size: block size + * @param x: key_cache.size(4) + * @param num_tokens: the number of tokens + * @param stream: CUDA Stream + */ +template +void vllm_reshape_and_cache_v4( + const T *key, + const T *value, + T *key_cache, + T *value_cache, + const int64_t *slot_mapping, + int key_stride, + int value_stride, + int num_heads, + int head_size, + int block_size, + int x, + int num_tokens, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param key: key. shape: [num_tokens, num_heads, head_size] + * @param value: value. shape: [num_tokens, num_heads, head_size] + * @param key_cache: key cache, shape: [num_blocks, num_heads, block_size, head_size] + * @param value_cache: value cache. shape: [num_blocks, num_heads, block_size, head_size] + * @param slot_mapping:The mapping position of the token in blocks. shape: [num_tokens] + * @param key_token_stride: key.stride(0) + * @param value_token_stride: value.stride(0) + * @param value_head_stride: value.stride(1) + * @param num_heads: the number of heads + * @param head_size: head size + * @param value_head_size: value head size, could be different from head size + * @param block_size: block size + * @param num_tokens: the number of tokens + * @param stream: CUDA Stream + */ +template +void vllm_reshape_and_cache( + const T *key, + const T *value, + T *key_cache, + T *value_cache, + const int64_t *slot_mapping, + int key_token_stride, + int value_token_stride, + int value_head_stride, + int num_heads, + int head_size, + int value_head_size, + int block_size, + int num_tokens, + cudaStream_t stream); + +/** + * @brief + * + * @param q_weight: quant weight + * @param aux_workspace: auxiliary workspace + * @param q_perm: g_idx + * @param height: q_weight.size(0) * 32 / bit + * @param width: q_weight.size(1) + * @param bit: quant weight bits + * @param stream: CUDA Stream + */ +template +void vllm_shuffle_exllama_weight( + T *q_weight, + T *aux_workspace, + const int *q_perm, + int height, + int width, + int bit, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 or float + * @tparam IS_NEOX: Determine whether to use Neox, that is, whether to use interleaved + * @param positions: token positions. shape: [num_tokens] + * @param query: query,shape: [num_tokens, num_heads * head_size] + * @param key: key, shape: [num_tokens, num_heads * head_size] + * @param cos_sin_cache: cos and sin value. shape: [max_position, head_size] + * @param rot_dim: cos_sin_cache.size(1) + * @param query_head_stride: stride on dim "head" + * @param query_token_stride: stride on dim "token" + * @param key_head_stride: stride on dim "head" + * @param key_token_stride: stride on dim "token" + * @param num_heads: the number of query heads + * @param num_kv_heads: the number of kv heads + * @param head_size: head size + * @param num_tokens: the number of tokens + * @param stream: CUDA Stream + */ + +template +void vllm_rotary_embedding(const int64_t *positions, + T *query, + T *key, + const T *cos_sin_cache, + int rot_dim, int query_head_stride, int query_token_stride, + int key_head_stride, int key_token_stride, + int num_heads, int num_kv_heads, int head_size, + int num_tokens, cudaStream_t stream); + + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 or float + * @tparam IS_NEOX: Determine whether to use Neox, that is, whether to use interleaved + * @param positions: token positions. shape: [num_tokens] + * @param query: query,shape: [num_tokens, num_heads * head_size] or [num_tokens, num_heads, head_size] + * @param key: key, shape: [num_tokens, num_kv_heads * head_size] or [num_tokens, num_kv_heads, head_size] + * @param cos_sin_cache: cos and sin value. shape: [max_position, head_size] + * @param scales: scales for key layer norm. shape: [head_size] + * @param bias: bias for key layer norm. shape: [head_size] + * @param key_out: result for saving key[nullptr will use inplace operation] + * @param rot_dim: cos_sin_cache.size(1) + * @param num_heads: the number of query heads + * @param num_kv_heads: the number of kv heads + * @param head_size: head size + * @param query_head_stride: stride of "num_heads" dim to support non contiguous query + * @param query_token_stride: stride of "num_tokens" dim to support non contiguous query + * @param key_head_stride: stride of "num_kv_heads" dim to support non contiguous key + * @param key_token_stride: stride of "num_tokens" dim to support non contiguous key + * @param eps: a value added to the denominator for numerical stability + * @param num_tokens: the number of tokens + * @param stream: CUDA Stream + */ +template +void vllm_rotary_embedding_with_key_layer_norm(const int64_t *positions, + scalar_t *query, + scalar_t *key, + const scalar_t *cos_sin_cache, + const scalar_t *scales, + const scalar_t *bias, + scalar_t *key_out, + int rot_dim, + int num_heads, + int num_kv_heads, + int head_size, + int64_t query_head_stride, + int64_t query_token_stride, + int64_t key_head_stride, + int64_t key_token_stride, + float eps, + int num_tokens, + cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 or float + * @tparam IS_NEOX: Determine whether to use Neox, that is, whether to use interleaved + * @param positions: token positions. shape: [num_tokens] + * @param query: query,shape: [num_tokens, num_heads * head_size] + * @param key: key, shape: [num_tokens, num_heads * head_size] + * @param cos_sin_cache: cos and sin value. shape: [max_position, head_size] + * @param cos_sin_cache_offsets: position offsets. shape: [num_tokens] + * @param rot_dim: cos_sin_cache.size(1) + * @param query_stride: query.stride(-2) + * @param key_stride: key.stride(-2) + * @param num_heads: the number of query heads + * @param num_kv_heads: the number of kv heads + * @param head_size: head size + * @param num_tokens: the number of tokens + * @param stream: CUDA Stream + */ +template +void vllm_batched_rotary_embedding(const int64_t *positions, + T *query, + T *key, + const T *cos_sin_cache, + const int64_t *cos_sin_cache_offsets, + int rot_dim, int query_stride, int key_stride, + int num_heads, int num_kv_heads, int head_size, + int num_tokens, cudaStream_t stream); + +/** + * @brief + * + * @param num_seqs: the number of sequences + * @param num_queries: NUM_QUERIES decode request numbers + * @param block_size: block size + * @param input_tokens: input token tensor + * @param sampled_token_ids: sampled token ids tensor + * @param input_positions: input positions tensor + * @param seq_lens: seq lens tensor + * @param slot_mapping: slot mapping tensor + * @param block_tables: block tables tensor + * @param block_tables_stride: block_tables.stride(0) + * @param stream: CUDA Stream + */ +void vllm_advance_step_flashattn(int num_seqs, int num_queries, int block_size, + long *input_tokens, + const long *sampled_token_ids, + long *input_positions, + int *seq_lens, + long *slot_mapping, + const int *block_tables, + long block_tables_stride, + cudaStream_t stream); + + +/** + * @brief + * + * @param positions: [num_tokens] + * @param long_prompt_offset: [num_tokens] + * @param long_short_cos_sin_cache: [num_tokens, head_dim] + * @param query: shape=[num_tokens, num_q_heads, head_dim] stride=[query_stride_0, query_stride_1, 1] + * @param key: shape=[num_tokens, num_kv_heads, head_dim] stride=[key_stride_0, key_stride_1, 1] + * @param out_query: shape=[num_tokens, num_q_heads, head_dim] stride=[out_query_stride_0, out_query_stride_1, 1] + * @param out_key: shape=[num_tokens, num_kv_heads, head_dim] stride=[out_key_stride_0, out_key_stride_1, 1] + */ + +template +void minicpm3_fused_rope( + const int64_t *positions, + const int64_t *long_prompt_offset, + const scalar_t *long_short_cos_sin_cache, + const scalar_t *query, + const scalar_t *key, + scalar_t *out_query, + scalar_t *out_key, + int64_t num_tokens, + int64_t num_q_heads, + int64_t num_kv_heads, + int64_t head_dim, + int64_t query_stride_0, + int64_t query_stride_1, + int64_t key_stride_0, + int64_t key_stride_1, + int64_t out_query_stride_0, + int64_t out_query_stride_1, + int64_t out_key_stride_0, + int64_t out_key_stride_1, + cudaStream_t stream); + +/** + * @brief + * + * @param k_nope: shape=(num_tokens, num_kv_heads, k_head_dim) stride=(k_nope_stride_0, k_nope_stride_1, 1) + * @param k_pe: shape=(num_tokens, 1, head_dim - k_head_dim) stride=(k_pe_stride_0, -1, 1) + * @param v: shape=(num_tokens, num_kv_heads, v_head_dim) stride=(v_stride_0, v_stride_1, 1) + * @param new_k: shape=(num_tokens, num_kv_heads, head_dim) contiguous + * @param new_v: shape=(num_tokens, num_kv_heads, head_dim) contiguous + */ +template +void minicpm3_fused_copy_kv( + const scalar_t *k_nope, + const scalar_t *k_pe, + const scalar_t *v, + scalar_t *new_k, + scalar_t *new_v, + int64_t num_tokens, + int64_t num_kv_heads, + int64_t head_dim, + int64_t k_head_dim, + int64_t v_head_dim, + int64_t k_nope_stride_0, + int64_t k_nope_stride_1, + int64_t k_pe_stride_0, + int64_t v_stride_0, + int64_t v_stride_1, + cudaStream_t stream); + + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 or float + * @param positions: token positions. shape: [num_tokens] + * @param query: query,shape: [num_tokens, num_heads * head_size] + * @param key: key, shape: [num_tokens, num_heads * head_size] + * @param cos_sin_cache: cos and sin value. shape: [max_position, head_size] + * @param offset: offset for lora, could be nullptr. shape: [max_position,] + * @param long_offset: add k or not. shape: [1, ] + * @param k: offset for long inputs + * @param rot_dim: cos_sin_cache.size(1) + * @param query_head_stride: stride on dim "head" + * @param query_token_stride: stride on dim "token" + * @param key_head_stride: stride on dim "head" + * @param key_token_stride: stride on dim "token" + * @param num_heads: the number of query heads + * @param num_kv_heads: the number of kv heads + * @param head_size: head size + * @param num_tokens: the number of tokens + * @param stream: CUDA Stream + */ +template +void vllm_rotary_embedding_phi(const int64_t *positions, + scalar_t *query, + scalar_t *key, + const scalar_t *cos_sin_cache, + const int64_t *offset, + const bool *long_offset, + const int64_t k, + int rot_dim, + int query_head_stride, + int query_token_stride, + int key_head_stride, + int key_token_stride, + int num_heads, + int num_kv_heads, + int head_size, + int num_tokens, + cudaStream_t stream); + + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 or float + * @param positions: token positions. shape: [num_tokens] + * @param query: query,shape: [num_tokens, num_heads * head_size] + * @param key: key, shape: [num_tokens, head_size] + * @param key_out: key_out, shape: [num_tokens, num_heads * head_size] + * @param cos_sin_cache: cos and sin value. shape: [max_position, head_size] + * @param offset: offset for lora, could be nullptr. shape: [max_position,] + * @param long_offset: add k or not. shape: [1, ] + * @param k: offset for long inputs + * @param rot_dim: cos_sin_cache.size(1) + * @param query_head_stride: stride on dim "head" + * @param query_token_stride: stride on dim "token" + * @param key_head_stride: stride on dim "head" + * @param key_token_stride: stride on dim "token" + * @param key_out_head_stride: stride on dim "head" + * @param key_out_token_stride: stride on dim "token" + * @param num_heads: the number of query heads + * @param head_size: head size + * @param num_tokens: the number of tokens + * @param stream: CUDA Stream + */ +template +void rotary_embedding_mla_phi(const int64_t *positions, + scalar_t *query, + scalar_t *key, + scalar_t *key_out, + const scalar_t *cos_sin_cache, + const int64_t *offset, + bool *long_offset, + int64_t k, + int rot_dim, + int query_head_stride, + int query_token_stride, + int key_head_stride, + int key_token_stride, + int key_out_head_stride, + int key_out_token_stride, + int num_heads, + int head_size, + int num_tokens, + cudaStream_t stream); + + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 or float + * @tparam IS_NEOX: Determine whether to use Neox, that is, whether to use interleaved + * @param positions: token positions. shape: [num_tokens] + * @param query: query,shape: [num_tokens, num_heads * head_size] + * @param key: key, shape: [num_tokens, num_heads * head_size] + * @param key_out: key_out, shape: [num_tokens, num_heads * head_size] + * @param cos_sin_cache: cos and sin value. shape: [max_position, head_size] + * @param offset: offset for lora, could be nullptr. shape: [max_position,] + * @param rot_dim: cos_sin_cache.size(1) + * @param query_head_stride: stride on dim "head" + * @param query_token_stride: stride on dim "token" + * @param key_head_stride: stride on dim "head" + * @param key_token_stride: stride on dim "token" + * @param key_out_head_stride: stride on dim "head" + * @param key_out_token_stride: stride on dim "token" + * @param num_heads: the number of query heads + * @param head_size: head size + * @param num_tokens: the number of tokens + * @param stream: CUDA Stream + */ +template +void rotary_embedding_mla(const int64_t *positions, + scalar_t *query, + scalar_t *key, + scalar_t *key_out, + const scalar_t *cos_sin_cache, + const int64_t *offset, + int rot_dim, + int query_head_stride, + int query_token_stride, + int key_head_stride, + int key_token_stride, + int key_out_head_stride, + int key_out_token_stride, + int num_heads, + int head_size, + int num_tokens, + cudaStream_t stream); + + +/** + * @brief + * + * @param key_nope: key_nope,shape: [num_tokens, num_heads, k_nope_dim] + * @param value_nope: value_nope,shape: [num_tokens, num_heads, v_head_dim] + * @param key: key, shape: [num_tokens, num_heads, head_size] + * @param value: value, shape: [num_tokens, num_heads, head_size] + * @param num_tokens: num_tokens + * @param num_heads: num_heads + * @param head_dim: head_size + * @param k_nope_dim: k_nope_dim + * @param v_head_dim: v_head_dim + * @param key_head_stride: stride on dim "head" + * @param key_token_stride: stride on dim "token" + * @param key_out_head_stride: stride on dim "head" + * @param key_out_token_stride: stride on dim "token" + * @param stream: CUDA Stream + */ +template +void copy_kv_mla( + const scalar_t *key_nope, + const scalar_t *value_nope, + scalar_t *key, + scalar_t *value, + int64_t num_tokens, + int64_t num_heads, + int64_t head_dim, + int64_t k_nope_dim, + int64_t v_head_dim, + int64_t k_nope_head_stride, + int64_t k_nope_token_stride, + int64_t v_nope_head_stride, + int64_t v_nope_token_stride, + cudaStream_t stream); + +/** + * @brief + * + * @param src_cache: src_cache,shape: [NUM_BLOCKS, BLOCK_SIZE,ENTRIES...] + * @param dst: workspace,shape: [TOT_TOKENS, ENTRIES...] + * @param block_table: block_table, shape: [BATCH, BLOCK_INDICES] + * @param cu_seq_lens: cu_seq_lens, shape: [BATCH+1] + * @param seq_starts: Optional: starting offsets per batch, shape: [BATCH] + * @param batch_size: batch size + * @param block_size: block size + * @param entry_size: entry size + * @param block_table_stride: stride on dim "BATCH" + * @param cache_block_stride: stride on dim "NUM_BLOCKS" + * @param cache_entry_stride: stride on dim "BLOCK_SIZE" + * @param dst_entry_stride: stride on dim "TOT_TOKENS" + * @param stream: CUDA Stream + */ +template +void vllm_gather_cache( + const scalar_t *src_cache, + scalar_t *dst, + const int32_t *block_table, + const int32_t *cu_seq_lens, + const int32_t *seq_starts, + const int64_t batch_size, + const int32_t block_size, + const int32_t entry_size, + const int64_t block_table_stride, + const int64_t cache_block_stride, + const int64_t cache_entry_stride, + const int64_t dst_entry_stride, + cudaStream_t stream); + +/** + * @brief + * + * @param src_cache: src_cache,shape: [NUM_BLOCKS, BLOCK_SIZE,ENTRIES...] + * @param src_cache_scale: src_cache,shape: [NUM_BLOCKS, BLOCK_SIZE,2] + * @param dst: workspace,shape: [TOT_TOKENS, ENTRIES...] + * @param block_table: block_table, shape: [BATCH, BLOCK_INDICES] + * @param cu_seq_lens: cu_seq_lens, shape: [BATCH+1] + * @param seq_starts: Optional: starting offsets per batch, shape: [BATCH] + * @param kv_lora_rank: kv_lora_rank + * @param batch_size: batch size + * @param block_size: block size + * @param entry_size: entry size + * @param block_table_stride: stride on dim "BATCH" + * @param cache_block_stride: stride on dim "NUM_BLOCKS" of src_cache + * @param scale_cache_block_stride: stride on dim "NUM_BLOCKS" of src_cache_scale + * @param cache_entry_stride: stride on dim "BLOCK_SIZE" of src_cache + * @param scale_cache_entry_stride: stride on dim "BLOCK_SIZE" of src_cache_scale + * @param dst_entry_stride: stride on dim "TOT_TOKENS" + * @param stream: CUDA Stream + */ +template +void vllm_gather_cache_int8( + const int8_t *src_cache, + const float *src_cache_scale, + scalar_t *dst, + const int32_t *block_table, + const int32_t *cu_seq_lens, + const int32_t *seq_starts, + const int64_t kv_lora_rank, + const int64_t batch_size, + const int32_t block_size, + const int32_t entry_size, + const int64_t block_table_stride, + const int64_t cache_block_stride, + const int64_t scale_cache_block_stride, + const int64_t cache_entry_stride, + const int64_t scale_cache_entry_stride, + const int64_t dst_entry_stride, + cudaStream_t stream); + +/** + * @brief + * + * @param kv_c: kv_c, shape: [num_tokens, kv_lora_rank] + * @param k_pe: query, shape: [num_tokens, 1(n), pe_dim] + * @param key_cache: key, shape: [num_tokens, block_size, (kv_lora_rank + pe_dim)] + * @param slot_mapping: slot_mapping, shape: [num_tokens] + * @param kv_lora_rank: kv_lora_rank + * @param pe_dim: pe_dim + * @param block_size: block_size + * @param kv_c_stride: stride on dim "num_tokens" + * @param k_pe_stride: stride on dim "num_tokens" + * @param block_stride: stride on dim "num_tokens" of key_cache + * @param dim_stride: stride on dim "block_size" of key_cache + * @param num_tokens: num_tokens + * @param stream: CUDA Stream + */ +template +void vllm_concat_and_cache_mla( + const scalar_t *kv_c, + const scalar_t *k_pe, + scalar_t *key_cache, + const int64_t *slot_mapping, + int kv_lora_rank, + int pe_dim, + int block_size, + int kv_c_stride, + int k_pe_stride, + int block_stride, + int dim_stride, + int num_tokens, + cudaStream_t stream); +/** + * @brief + * + * @param kv_c: kv_c, shape: [num_tokens, kv_lora_rank] + * @param kv_c_scale: kv_c_scale, shape: [num_tokens] + * @param k_pe: query, shape: [num_tokens, 1(n), pe_dim] + * @param k_pe_scale: query, shape: [num_tokens, 1(n)] + * @param key_cache: key, shape: [num_tokens, block_size, (kv_lora_rank + pe_dim)] + * @param key_cache_scale: key, shape: [num_tokens, block_size, 2] + * @param slot_mapping: slot_mapping, shape: [num_tokens] + * @param kv_lora_rank: kv_lora_rank + * @param pe_dim: pe_dim + * @param block_size: block_size + * @param kv_c_stride: stride on dim "num_tokens" + * @param kv_c_scale_stride: stride on dim "num_tokens" + * @param k_pe_stride: stride on dim "num_tokens" + * @param k_pe_scale_stride: stride on dim "num_tokens" + * @param block_stride: stride on dim "num_tokens" of key_cache + * @param scale_block_stride: stride on dim "num_tokens" of key_cache_scale + * @param dim_stride: stride on dim "block_size" of key_cache + * @param scale_dim_stride: stride on dim "block_size" of key_cache_scale + * @param num_tokens: num_tokens + * @param stream: CUDA Stream +*/ +template +void vllm_concat_and_cache_mla_int8( + const scalar_t *kv_c, + const float *kv_c_scale, + const scalar_t *k_pe, + const float *k_pe_scale, + scalar_t *key_cache, + float *key_cache_scale, + const int64_t *slot_mapping, + int kv_lora_rank, + int pe_dim, + int block_size, + int kv_c_stride, + int kv_c_scale_stride, + int k_pe_stride, + int k_pe_scale_stride, + int block_stride, + int scale_block_stride, + int dim_stride, + int scale_dim_stride, + int num_tokens, + cudaStream_t stream); + +/** + * @brief + * + * @param output, shape: [seq_len, num_heads, head_dim] + * @param output_lse, shape: [num_heads, seq_len] + * @param prefix_output, shape: [seq_len, num_heads, head_dim] + * @param prefix_lse, shape: [num_heads, seq_len] + * @param suffix_output, shape: [seq_len, num_heads, head_dim] + * @param suffix_lse, shape: [num_heads, seq_len] + */ +template +void merge_attn_states( + scalar_t *output, + float *output_lse, + const scalar_t *prefix_output, + const float *prefix_lse, + const scalar_t *suffix_output, + const float *suffix_lse, + int num_heads, + int seq_len, + int head_dim, + cudaStream_t stream); + + + +/* + MARLIN_FORMAT_K16N32 + w:(batch, k/16, n/32, 64) int32 pack order:[0 2 4 6 1 3 5 7] + s:(batch, k_groups, n/32, 32) float16 32 data order:[0 16 1 17 ... 15 31] + z:(batch, k_groups, n/32, 32) int4 32 data order:[0 16 1 17 ... 15 31] + MARLIN_FORMAT_K16N32_GROUPED_ON_N + w:(batch, k/16, n/32, 64) int32 pack order:[0 2 4 6 1 3 5 7] + s:(batch, n_groups, k) float16 + z:(batch, n_groups, k/8) int4 pack order:[0 1 2 3 4 5 6 7] + MARLIN_FORMAT_K16N16 + w:(batch, k/16, n/16, 64) int32 pack order:[0 1 2 3] + s:(batch, k_groups, n) float32 + MARLIN_FORMAT_K16N16_GROUPED_ON_N + w:(batch, k/16, n/16, 64) int32 pack order:[0 1 2 3] + s:(batch, n_groups, k) float32 +*/ +typedef enum { + MARLIN_FORMAT_K16N32, + MARLIN_FORMAT_K16N32_GROUPED_ON_N, + MARLIN_FORMAT_K16N16, + MARLIN_FORMAT_K16N16_GROUPED_ON_N, +} MarlinFormat; + +/* + ORIGIN_FORMAT_AWQ, + pack_order:[0 2 4 6 1 3 5 7] + w:(batch, k, n/8) int32 + s:(batch, k_groups, n) float16 + z:(batch, k_groups, n/8) int32 + ORIGIN_FORMAT_GPTQ, + pack_order:[0 1 2 3 4 5 6 7] + w:(batch, k/8, n) int32 + s:(batch, k_groups, n) float16 + z:(batch, k_groups, n/8) int32 + ORIGIN_FORMAT_GPTQ_GROUPED_N, + pack_order:[0 2 4 6 1 3 5 7] + w:(batch, k/8, n) int32 + s:(batch, n_groups, k) float16 + z:(batch, n_groups, k/8) int32 + ORIGIN_FORMAT_INT8 + w:(batch, k, n) int8 +*/ +typedef enum { + ORIGIN_FORMAT_AWQ, + ORIGIN_FORMAT_GPTQ, + ORIGIN_FORMAT_GPTQ_GROUPED_N, + ORIGIN_FORMAT_INT8, +} WeightFormat; + +typedef enum { + PACK_ORDER_01234567, + PACK_ORDER_02461357, +} PackOrder; + +/** + * @brief + * + * @tparam DType: input type, half or bfloat16 + * @param input: input tensor, shape: batch_first ? [batch_count, m, k] : [m, batch_count, k] + * @param weight: marlin repack weights, shape: [batch_count, k/16, n/32, 64] + * @param scale: marlin repack scale, shape: weight_format == "k16n32" ? [batch, k_groups, n] : [batch, n_groups, k] + * @param zero: marlin repack zero, shape: weight_format == "k16n32" ? [batch, k_groups, n/8] : [batch, n_groups, k/8] + * @param bias: bias for result, TODO + * @param out: output tensor, shape: batch_first ? [batch_count, m, n] : [m, batch_count, n] + * @param aux: workspace for kernel + * @param batch_count: batched gemm paraments + * @param m: gemm paraments + * @param k: gemm paraments + * @param n: gemm paraments + * @param group_size: group size of quant + * @param pad_k: stride for k dimension of input + * @param batch_first: describe format of input and output + * @param weight_format: describe format of weight + * @param stream: CUDA Stream + */ +template +void marlin_w4a16(const DType *input, const int32_t *weight, const DType *scale, const int32_t *zero, const DType *bias, + DType *out, float *aux, int batch_count, int m, int k, int n, int group_size, int pad_k, bool batch_first, MarlinFormat weight_format, cudaStream_t stream); + + +/** + * @brief + * + * @param weight: origin weight tensor + * @param repack_weight: marlin repack weight tensor + * @param scale: origin scale tensor + * @param repack_scale: marlin repack scale tensor + * @param zero: origin zero tensor + * @param repack_zero: marlin repack zero tensor + * @param batch_count: batched gemm paraments + * @param n: gemm paraments + * @param k: gemm paraments + * @param groups: groups of quant + * @param origin_format: describe format of origin weight + * @param origin_pack_order: describe pack order of origin weight + * @param marlin_format: describe format of repack weight + * @param stream: CUDA Stream + */ +void marlin_w4_weight_repack(const void *weight, void *repack_weight, + const void *scale, void *repack_scale, + const void *zero, void *repack_zero, + int batch_count, int n, int k, int groups, + WeightFormat origin_format, PackOrder origin_pack_order, MarlinFormat marlin_format, cudaStream_t stream); + +/** + * @brief + * + * @tparam DType: input type, half or bfloat16 + * @param input: input tensor, shape: batch_first ? [batch_count, m, k] : [m, batch_count, k] + * @param weight: marlin repack weights, shape: [batch_count, k/16, n/16, 64] + * @param scale: marlin repack scale, shape: weight_format == "k16n16" ? [batch, k_groups, n] : [batch, n_groups, k] + * @param bias: bias for result, TODO + * @param out: output tensor, shape: batch_first ? [batch_count, m, n] : [m, batch_count, n] + * @param aux: workspace for kernel + * @param batch_count: batched gemm paraments + * @param m: gemm paraments + * @param k: gemm paraments + * @param n: gemm paraments + * @param group_size: group size of quant + * @param pad_k: stride for k dimension of input + * @param batch_first: describe format of input and output + * @param weight_format: describe format of weight + * @param stream: CUDA Stream + */ +template +void marlin_w8a16(const DType *input, const int32_t *weight, const float *scale, const DType *bias, + DType *out, float *aux, int batch_count, int m, int k, int n, int group_size, int pad_k, bool batch_first, MarlinFormat weight_format, cudaStream_t stream); + +/** + * @brief + * + * @param weight: origin weight tensor + * @param repack_weight: marlin repack weight tensor + * @param scale: origin scale tensor + * @param repack_scale: marlin repack scale tensor + * @param batch_count: batched gemm paraments + * @param n: gemm paraments + * @param k: gemm paraments + * @param groups: groups of quant + * @param origin_format: describe format of origin weight + * @param marlin_format: describe format of repack weight + * @param stream: CUDA Stream + */ +void marlin_w8_weight_repack(const void *weight, void *repack_weight, + const void *scale, void *repack_scale, + int batch_count, int n, int k, int groups, + WeightFormat origin_format, MarlinFormat marlin_format, cudaStream_t stream); + +// ======================================================== +// bert unpad +// ======================================================== +/** + * @brief bert layernorm fused add residual + * + * @tparam T + * @param input: shape:[num_tokens, hidden_size] half bf16 + * @param residual: shape:[num_tokens, hidden_size] same as input + * @param ln_weight： layernorm weight,shape:[hidden_size] same as input + * @param ln_bias：layernorm bias,shape:[hidden_size] same as input + * @param output: shape:[num_tokens, hidden_size] same as input + * @param num_tokens: int ,total tokens in a batch + * @param hidden_size: HiddenSize + * @param epsilon: float + * @param stream: CUDA Stream + */ +template +void bert_add_norm(const T *input, const T *residual, + const T *ln_weight, const T *ln_bias, + T *output, int num_tokens, int hidden_size, + float epsilon, cudaStream_t stream); +/** + * @brief bert embeding same as transformers + * + * @tparam T + * @tparam TYPE_INT: type for token_ids pos_ids type_ids + * @param token_weight: shape: [vocab_size, hidden_size] half bf16 + * @param pos_weight: shape: [pos_size, hidden_size] same as token_weight + * @param type_weight: shape:[type_size, hidden_size] same as token_weight + * @param ln_weight: layernorm weight,shape:[hidden_size] same as token_weight + * @param ln_bias: layernorm bias,shape:[hidden_size] same as token_weight + * @param token_ids: shape: [num_tokens] + * @param pos_ids: shape: [num_tokens] + * @param type_ids: shape: [num_tokens] + * @param output: shape: [num_tokens, hidden_size] + * @param num_tokens: int ,total tokens in a batch + * @param hidden_size: HiddenSize + * @param epsilon: float + * @param stream: CUDA Stream + */ +template +void bert_embedding(const T *token_weight, const T *pos_weight, + const T *type_weight, const T *ln_weight, + const T *ln_bias, + const TYPE_INT *token_ids, const TYPE_INT *, + const TYPE_INT *type_ids, + T *output, int num_tokens, int hidden_size, + float epsilon, cudaStream_t stream); + +/** + * @brief bert output numtokens unpack to batch,tokens + * + * @tparam T + * @tparam TYPE_INT + * @param logits: shape:[num_tokens, 2] half bf16 + * @param cu_seq_len: shape:[batch+1],same as in flash atten,accumlate seq_len in a batch,first is 0 + * @param start_logits: shape:[ batch, max_seq_len] half bf16 + * @param end_logits: shape:[ batch, max_seq_len] half bf16 + * @param batch： Batch Size + * @param max_seq_len + * @param stream + */ +template +void bert_unpack_start_end_logits(const T *logits, const TYPE_INT *cu_seq_len, + T *start_logits, T *end_logits, + int batch, int max_seq_len, + cudaStream_t stream); + +// ======================================================== +// Linalg.solve +// ======================================================== +/** + * @brief + * + * @tparam T: input type, float + * @param A: tensor of shape [*, n, n] where * is zero or more batch dimensions. + * @param B: right-hand side tensor of shape [*, n] or [*, n, k] or or [*, k, n], + * where * is zero or more batch dimensions. + * @param X: output tensor, shape: [*, n] or [*, n, k] or [*, k,n] + * @param batch: batch size, the value is A.numel() / (n * n) + * @param n: One of the dimensions of the param B tensor + * @param k: One of the dimensions of the param B tensor + * @param stream: CUDA Stream + */ +template +void gauss_small(const T *A, const T *B, T *X, int batch, int n, int k, cudaStream_t stream); + +// ======================================================== +// store_kv_cache +// ======================================================== + +/** + * @brief + * + * @tparam T: input type, half, bfloat16 + * @param k: key. shape: [batch_size, seqlen_new, head_num, head_dim] + * @param v: value. shape: [batch_size, seqlen_new, head_num, head_dim] + * @param k_cache: key cache. shape: [batch_size_cache, seqlen_cache, head_num, head_dim] + * @param v_cache: value cache. shape: [batch_size_cache, seqlen_cache, head_num, head_dim] + * @param cache_batch_idx: The indices used to index into the KV cache. shape: [batch_size,] + * @param cache_seqlens: The sequence lengths of the KV cache. shape: [batch_size,] + * @param k_stride_1: k.stride(0) + * @param k_stride_2: k.stride(1) + * @param k_stride_3: k.stride(1) + * @param v_stride_1: v.stride(0) + * @param v_stride_2: v.stride(1) + * @param v_stride_3: v.stride(1) + * @param batch_size: k.size(0) + * @param seq_len_new: k.size(1) + * @param seqlen_cache: k_cache.size(1) + * @param head_num: k.size(2) + * @param head_dim: k.size(3) + * @param stream: CUDA Stream + */ +template +void store_kv_cache(const T *k, const T *v, + T *k_cache, T *v_cache, + const int32_t *cache_batch_idx, const int32_t *cache_seqlens, + int64_t k_stride_1, int64_t k_stride_2, int64_t k_stride_3, + int64_t v_stride_1, int64_t v_stride_2, int64_t v_stride_3, + int batch_size, int seq_len_new, + int seqlen_cache, + int head_num, int head_dim, + cudaStream_t stream); + +// ======================================================== +// T5 model +// ======================================================== + +/** + * @brief t5_split_qkv + * + * @tparam T: input type, half or bfloat16 + * @param qkv: input tensor, shape: [batch_size, seq_len, hidden_size*3] + * @param q: query tensor, shape: [batch_size, head_num, seq_len, head_dim] + * @param k: key tensor, shape: [batch_size, head_num, seq_len, head_dim] + * @param v: value tensor, shape: [batch_size, head_num, seq_len, head_dim] + * @param batch: int, batch size + * @param seq_len: int, seq_len + * @param head_num: int, the number of head + * @param head_dim: int, head dim + * @param stream: CUDA Stream + */ +template +void t5_split_qkv(const T *qkv, T *q, T *k, T *v, int batch, + int seq_len, int head_num, int head_dim, cudaStream_t stream); + +/** + * @brief + * + * @tparam T: input type, half or bfloat16 + * @param qkv: input tensor, shape: [batch_size,1,hidden_size*3],hidden_size = head_num*head_dim + * @param past_key: past key tensor, shape: [batch_size, head_num, seq_len-1, head_dim] + * @param past_value: past value tensor, shape: [batch_size, head_num, seq_len-1, head_dim] + * @param q: query tensor, shape: [batch_size, head_num, 1, head_dim] + * @param k: key tensor, shape: [batch_size, head_num, 1, head_dim] + * @param v: value tensor, shape: [batch_size, head_num, 1, head_dim] + * @param batch: int, batch size + * @param seq_len: int, seq_len + * @param head_num: int, the number of head + * @param head_dim: int, head dim + * @param stream: CUDA Stream + */ +template +void t5_split_qkv_update_kv_cache(const T *qkv, const T *past_key, const T *past_value, + T *q, T *k, T *v, int batch, + int seq_len, int head_num, int head_dim, + cudaStream_t stream); + +}// namespace ixformer::kernels::infer diff --git a/ixformer_sdk/csrc/include/ixformer/kernels/status.h b/ixformer_sdk/csrc/include/ixformer/kernels/status.h new file mode 100644 index 00000000..495a4d84 --- /dev/null +++ b/ixformer_sdk/csrc/include/ixformer/kernels/status.h @@ -0,0 +1,20 @@ +#pragma once + +#include + +namespace ixformer::kernels { + +enum KernelStatus { + kernelSuccess, + kernelFail, + kernelCudaError, + kernelInvalidArgument, + kernelCuinferError, + kernelUnsupported, +}; + + +std::string to_string(KernelStatus status); + + +}// namespace ixformer::kernels diff --git a/ixformer_sdk/csrc/include/ixformer/kernels/tensor.h b/ixformer_sdk/csrc/include/ixformer/kernels/tensor.h new file mode 100644 index 00000000..a19b84fb --- /dev/null +++ b/ixformer_sdk/csrc/include/ixformer/kernels/tensor.h @@ -0,0 +1,92 @@ +#pragma once + +#include + +namespace ixformer::kernels { + +const uint8_t MAX_TENSOR_NDIM = 8; + +// align with at::ScalarType +enum DType { + Byte = 0, + Char = 1, + Short = 2, + Int = 3, + Long = 4, + Half = 5, + Float = 6, + Double = 7, + ComplexHalf = 8, + ComplexFloat = 9, + ComplexDoubl = 10, + Bool = 11, + QInt8 = 12, + QUInt8 = 13, + QInt32 = 14, + BFloat16 = 15, + QUInt4x2 = 16, + QUInt2x4 = 17, + Bits1x8 = 18, + Bits2x4 = 19, + Bits4x2 = 20, + Bits8 = 21, + Bits16 = 22, + Float8_e5m2 = 23, + Float8_e4m3fn = 24, + Undefined = 25, + NumOptions = 26 +}; + +struct TensorDesc { + +public: + // delete default constructor + TensorDesc() = delete; + // All information must be (should be) prepared when constructing a TensorDesc object. + TensorDesc(DType scalar_type, void *data_ptr, int64_t numel, int64_t dim, const int64_t *size, const int64_t *stride, bool is_contiguous, bool is_cuda) + : dtype(scalar_type), ptr(data_ptr), nnumel(numel), ndim(dim), sizes(size), strides(stride), contiguous(is_contiguous), cuda(is_cuda) {} + + inline DType scalar_type() const { + return dtype; + } + + inline void *data_ptr() const { + return ptr; + } + + inline int64_t numel() const { + return nnumel; + } + + inline int64_t dim() const { + return ndim; + } + + inline int64_t size(int64_t dim) const { + return dim < 0 ? sizes[ndim - dim] : sizes[dim]; + } + + inline int64_t stride(int64_t dim) const { + return dim < 0 ? strides[ndim - dim] : strides[dim]; + } + + inline bool is_contiguous() const { + return contiguous; + } + + inline bool is_cuda() const { + return cuda; + } + +private: + void *ptr{nullptr}; + DType dtype; + int64_t nnumel{0}; + int64_t ndim{0}; + const int64_t *sizes{nullptr}; + const int64_t *strides{nullptr}; + bool contiguous{false}; + bool cuda{false}; +}; + +}// namespace ixformer::kernels diff --git a/ixformer_sdk/distributed/__init__.py b/ixformer_sdk/distributed/__init__.py new file mode 100644 index 00000000..f57c7252 --- /dev/null +++ b/ixformer_sdk/distributed/__init__.py @@ -0,0 +1 @@ +from ._distributed import * diff --git a/ixformer_sdk/distributed/_distributed.py b/ixformer_sdk/distributed/_distributed.py new file mode 100644 index 00000000..1c482f81 --- /dev/null +++ b/ixformer_sdk/distributed/_distributed.py @@ -0,0 +1,481 @@ +import warnings +from collections import defaultdict +from typing import List, Optional, Tuple + +import torch +import torch.distributed as dist +import torch.distributed.distributed_c10d as c10d +from ixformer._C import _distributed as cdist +from ixformer._C._distributed import comm +from ixformer._C._distributed.comm import ( + AllGatherAlgo, + AllReduceAlgo, + BroadcastAlgo, + ReduceAlgo, + ReduceOp, + ReduceScatterAlgo, + SendAlgo, +) +from ixformer.core.multi_level_cache import MultiLevelCache +from torch import Tensor +from torch.distributed import ProcessGroup + +from ixformer.core import config + +IxformerCommType = int +RecvAlgo = SendAlgo + +_GROUP_TO_IXFC_COMM_CACHE = MultiLevelCache() +_IXFC_COMM_TO_GROUP_CACHE = MultiLevelCache() + + +def get_store(group: dist.ProcessGroup = None) -> dist.Store: + if group is None: + group = c10d._get_default_group() + + return c10d._pg_map[group][1] + + +class StoreWrapper(cdist.comm.C10dStoreWrapper): + _GROUP_COUNT = defaultdict(dict) + + def __init__(self, group: ProcessGroup): + super().__init__() + self.store = get_store() + + ranks = dist.get_process_group_ranks(group) + + group_key = "_".join([str(r) for r in ranks]) + if group not in self._GROUP_COUNT[group_key]: + self._GROUP_COUNT[group_key][group] = len(self._GROUP_COUNT[group_key]) + group_count = self._GROUP_COUNT[group_key][group] + + self.prefix = f"gid_{group_count}_" + group_key + + def _gen_unique_key(self, key): + return f"{self.prefix}_{key}" + + def set(self, key: str, value: str): + key = self._gen_unique_key(key) + self.store.set(key, value) + + def get(self, key: str) -> str: + key = self._gen_unique_key(key) + self.store.wait([key]) + return self.store.get(key).decode("utf8") + + +def init_comm_with_store(group=None, shmsize: int = None): + if group is None: + group = c10d._get_default_group() + + world_size = dist.get_world_size(group=group) + rank = dist.get_group_rank(group=group, global_rank=dist.get_rank()) + + if shmsize is None: + shmsize = config.IXFORMER_COMM_SHM_SIZE + + store_wrapper = StoreWrapper(group=group) + ixfc_comm = cdist.comm.init_communicator_by_store( + store=store_wrapper, world_size=world_size, rank=rank, max_shm_mem_size=shmsize + ) + + _GROUP_TO_IXFC_COMM_CACHE.set(group, ixfc_comm) + _IXFC_COMM_TO_GROUP_CACHE.set(ixfc_comm, group) + return ixfc_comm + + +_sub_store = None + + +def create_nccl_unique_id(addr: str, port: str, world_size: int, rank: int): + global _sub_store + _sub_store = dist.TCPStore( + host_name=addr, port=int(port), world_size=world_size, is_master=rank == 0 + ) + store_key = "ncclUniqueId" + if rank == 0: + commid = cdist.comm.create_nccl_unique_id() + _sub_store.set(store_key, commid) + else: + _sub_store.wait([store_key]) + commid = _sub_store.get(store_key).decode("utf8") + + return commid + + +def init_comm_with_eth( + addr: str, port: str, world_size: int, rank: int, shmsize: int = None +): + commid = create_nccl_unique_id(addr, port, world_size=world_size, rank=rank) + return cdist.comm.init_communicator_by_nccl_id(commid, world_size, rank, shmsize) + + +def _check_group(group: Optional[ProcessGroup] = None): + if group is None: + group = c10d._get_default_group() + + if isinstance(group, ProcessGroup): + ixfc_comm = _GROUP_TO_IXFC_COMM_CACHE.get(group, None) + if ixfc_comm is None: + return init_comm_with_store(group) + return ixfc_comm + + return group + + +def get_comm_group_stream(group: Optional[ProcessGroup] = None): + group = _check_group(group) + return comm.get_comm_group_stream(group) + + +def set_comm_group_stream(stream: int, group: Optional[ProcessGroup] = None): + group = _check_group(group) + return comm.set_comm_group_stream(group, stream) + + +def get_group_rank(group: Optional[ProcessGroup], global_rank) -> int: + """将 global rank 映射到 group 中的相对 rank""" + if isinstance(group, IxformerCommType): + _pg = _IXFC_COMM_TO_GROUP_CACHE.get(group, None) + if _pg is None: + return global_rank + else: + group = _IXFC_COMM_TO_GROUP_CACHE.get(group) + + if group is None: + group = c10d._get_default_group() + + return dist.get_group_rank(group, global_rank) + + +def get_global_rank(group: Optional[ProcessGroup], group_rank: int) -> int: + """将一个 group rank 映射到 global rank""" + if group is None: + group = c10d._get_default_group() + return c10d.get_global_rank(group, group_rank) + + +def get_process_group_ranks(group: Optional[ProcessGroup] = None) -> List[int]: + """获取 Group 的 global ranks""" + if group is None: + group = c10d._get_default_group() + return c10d.get_process_group_ranks(group) + + +def new_group(ranks: List[int] = None, shmsize=None, *args, **kwargs): + """通过 global ranks 去创建一个通讯组""" + group = c10d.new_group(ranks, *args, **kwargs) + + if ranks is None: + ranks = dist.get_process_group_ranks(group) + + if get_rank() in ranks: + init_comm_with_store(group=group, shmsize=shmsize) + return group + + +def new_subgroups_by_enumeration( + ranks_per_subgroup_list, shmsize=None, *args, **kwargs +) -> Tuple[ProcessGroup, List[ProcessGroup]]: + """ + 通过一组 global ranks 去创建通讯组 + + :param ranks_per_subgroup_list: global ranks + :return: 返回当前 rank 所在的通讯组 和 新的 subgroups + """ + self_group, other_group = c10d.new_subgroups_by_enumeration( + ranks_per_subgroup_list, *args, **kwargs + ) + init_comm_with_store(self_group, shmsize=shmsize) + return self_group, other_group + + +def destroy_process_group(group: Optional[ProcessGroup] = None): + """销毁 Group""" + if group is None: + group = c10d._get_default_group() + ixfc_comm = _GROUP_TO_IXFC_COMM_CACHE.get(group, None) + + if ixfc_comm is None: + dist.destroy_process_group(group) + else: + comm.destroy(ixfc_comm) + dist.destroy_process_group(group) + + +def get_rank(group: Optional[ProcessGroup] = None) -> int: + """获取当前进程的 Rank，如果 group 是 null，那么返回的是 Global Rank, 否则返回的相对的 Rank，即在当前组中的 rank""" + return c10d.get_rank(group) + + +def get_world_size(group: Optional[ProcessGroup] = None) -> int: + """获取 Group 中的成员大小""" + return c10d.get_world_size(group) + + +def barrier(group: Optional[ProcessGroup] = None, use_comm_stream: bool = False): + """同步 Group 中的 rank""" + group = _check_group(group) + comm.barrier(group, use_comm_stream) + + +def isend( + tensor: Tensor, + dst: int, + group: Optional[ProcessGroup] = None, + use_comm_stream: bool = False, +): + dst = get_group_rank(group, dst) + group = _check_group(group) + return comm.send(group, tensor, dst, use_comm_stream, SendAlgo.kNone) + + +def send(*args, **kwargs): + warnings.warn("not support sync mode, as async to call.") + return isend(*args, **kwargs) + + +def irecv( + tensor: torch.Tensor, + src: int, + group: Optional[ProcessGroup] = None, + use_comm_stream: bool = False, +): + src = get_group_rank(group, src) + group = _check_group(group) + return comm.recv(group, tensor, src, use_comm_stream, SendAlgo.kNone) + + +def recv(*args, **kwargs): + warnings.warn("not support sync mode, as async to call.") + return irecv(*args, **kwargs) + + +def point_to_point( + tensor: Tensor, + src: int, + dst: int, + group: Optional[ProcessGroup] = None, + use_comm_stream: bool = False, +): + """在 src rank 发送 tensor，在 dst_rank 上接收数据到 tensor 中""" + src = get_group_rank(group, src) + dst = get_group_rank(group, dst) + group = _check_group(group) + return comm.p2p(group, tensor, src, dst, use_comm_stream) + + +def reduce( + tensor, + root: int, + op=ReduceOp.SUM, + group: Optional[ProcessGroup] = None, + async_op=False, + out: Tensor = None, + use_comm_stream: bool = False, +): + """ + Example: + ixf_tensor = torch.tensor([1], device="cuda") + ixfd.reduce(ixf_tensor, 1, async_op=True) + print("rank {rank}:", ixf_tensor) + + # output + rank 0: tensor([1], device='cuda:0') + rank 1: tensor([4], device='cuda:1') + rank 2: tensor([1], device='cuda:2') + rank 3: tensor([1], device='cuda:3') + """ + + if not async_op: + raise RuntimeError("Not support sync operation now.") + + if out is None: + out = tensor + + root = get_group_rank(group, root) + group = _check_group(group) + return comm.reduce(group, tensor, out, op, root, use_comm_stream, ReduceAlgo.kNone) + + +def broadcast( + tensor: Tensor, + src: int, + group: Optional[ProcessGroup] = None, + async_op=False, + out: Tensor = None, + use_comm_stream: bool = False, +): + """ + Example: + ixf_tensor = torch.tensor([rank], device="cuda") + ixfd.broadcast(ixf_tensor, 1, async_op=True) + print("rank {rank}: ", ixf_tensor) + + # output + rank 0: tensor([1], device='cuda:0') + rank 1: tensor([1], device='cuda:1') + rank 2: tensor([1], device='cuda:2') + rank 3: tensor([1], device='cuda:3') + """ + if not async_op: + raise RuntimeError("Not support sync operation now.") + + if out is None: + out = tensor + + src = get_group_rank(group, src) + group = _check_group(group) + return comm.broadcast(group, tensor, out, src, use_comm_stream, BroadcastAlgo.kNone) + + +def reduce_scatter_tensor( + output: Tensor, + input: Tensor, + op=ReduceOp.SUM, + group: Optional[ProcessGroup] = None, + async_op=False, + use_comm_stream: bool = False, +): + """ + Example: + ixf_tensor_out = torch.zeros(2, dtype=torch.int64, device="cuda") + tensor_in = torch.arange(world_size * 2, dtype=torch.int64, device="cuda") + # tensor_in: tensor([0, 1, 2, 3, 4, 5, 6, 7], device='cuda:0') + + ixfd.reduce_scatter_tensor(ixf_tensor_out, tensor_in, async_op=True) + print("rank {rank}:", ixf_tensor_out) + + # output + rank 0: tensor([0, 4], device='cuda:0') + rank 1: tensor([ 8, 12], device='cuda:1') + rank 2: tensor([16, 20], device='cuda:2') + rank 3: tensor([24, 28], device='cuda:3') + """ + if not async_op: + raise RuntimeError("Not support sync operation now.") + + group = _check_group(group) + return comm.reduce_scatter( + group, input, output, op, use_comm_stream, ReduceScatterAlgo.kNone + ) + + +def all_reduce( + tensor: Tensor, + op=ReduceOp.SUM, + group: Optional[ProcessGroup] = None, + async_op=False, + out: Tensor = None, + algo: AllReduceAlgo = AllReduceAlgo.kNone, + use_comm_stream: bool = False, +): + """ + Args: + tensor: inpute tensor + op: ReduceOp: SUM, MIN or MAX + group: communicator group + async_op: ixformer support async mode + out: output tensor + algo: AllReduce Algo: Auto, Quant, QuantL1, QuantL2, NCCL, Ring, AllGatherSum, BroadcastSum + use_comm_stream: ixformer support set communication stream by ixformer.distributed.set_comm_group_stream, + if true, submit the kernels of communication to communication stream, + if false, use current stream by torch.cuda.current_stream + Returns: out + + Example: + >>> # All tensors below are of torch.int64 type. + >>> # We have 2 process groups, 2 ranks. + >>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank + >>> tensor + tensor([1, 2]) # Rank 0 + tensor([3, 4]) # Rank 1 + >>> ixfd.all_reduce(tensor, op=ReduceOp.SUM, async_op=True) + >>> tensor + tensor([4, 6]) # Rank 0 + tensor([4, 6]) # Rank 1 + """ + if not async_op: + raise RuntimeError("Not support sync operation now.") + + group = _check_group(group) + + if out is None: + out = tensor + + comm.all_reduce( + group, + tensor, + out, + op, + use_comm_stream=use_comm_stream, + algo=algo, + ) + + +def all_gather_into_tensor( + output: Tensor, + input: Tensor, + group: Optional[ProcessGroup] = None, + async_op=False, + use_comm_stream: bool = False, +): + """ + Example: + tensor_in = torch.arange(2, dtype=torch.int64, device="cuda") + 1 + 2 * rank + rank 0: tensor in: tensor([1, 2], device='cuda:0') + rank 1: tensor in: tensor([3, 4], device='cuda:1') + rank 2: tensor in: tensor([5, 6], device='cuda:2') + rank 3: tensor in: tensor([7, 8], device='cuda:3') + + ixf_tensor_out = torch.zeros(world_size * 2, dtype=torch.int64, device="cuda") + ixfd.all_gather_into_tensor(ixf_tensor_out, tensor_in, async_op=True) + print("rank {rank}:", ixf_tensor_out) + + # output: + rank 0: tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:0') + rank 1: tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:1') + rank 2: tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:2') + rank 3: tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:3') + """ + if not async_op: + raise RuntimeError("Not support sync operation now.") + + group = _check_group(group) + return comm.all_gather( + group, input, output, use_comm_stream, algo=AllGatherAlgo.kNone + ) + + +def gather( + tensor, + gather_list=None, + dst=0, + group: Optional[ProcessGroup] = None, + async_op=False, + use_comm_stream: bool = False, +): + """ + Example: + >>> # We have 2 process groups, 2 ranks. + >>> tensor = torch.tensor(rank+1,dtype=torch.float32).cuda() + >>> tensor + tensor(1.) # Rank 0 + tensor(2.) # Rank 1 + >>> gather_list = [torch.zeros(1).cuda() for _ in range(rank)] if rank == dst else None + >>> gather_list + [tensor([0,]),tensor([1,])] # Rank 0 + None # Rank 1 + ixfd.gather(tensor,gather_list,0,async_op=True) + >>> gather_list + [tensor([1.]),tensor([2.])] # Rank 0 + None # Rank 1 + """ + gather_list = gather_list if gather_list is not None else [] + if not async_op: + raise RuntimeError("Not support sync operation now.") + + dst = get_group_rank(group, dst) + group = _check_group(group) + return comm.gather(group, tensor, gather_list, dst, use_comm_stream) diff --git a/ixformer_sdk/distributed/overlap_comm.py b/ixformer_sdk/distributed/overlap_comm.py new file mode 100644 index 00000000..833e0874 --- /dev/null +++ b/ixformer_sdk/distributed/overlap_comm.py @@ -0,0 +1,412 @@ +import abc +import enum +import os +from contextlib import contextmanager, nullcontext +from typing import List, Optional + +import torch.cuda +from ixformer.core.dispatcher import Dispatcher + +from ixformer.core import config + +from . import _distributed as ixfd + + +class SplitOverlapComm(Dispatcher): + def __init__(self, num_chunks, num_compute_streams=None, comm_group=None): + """ + Args: + num_chunks: the number of chunks + num_compute_streams: the number of compute streams, default: 1 + comm_group: communicator group + """ + + self._num_chunks = num_chunks + self._num_compute_streams = num_compute_streams or 1 + self._comm_group = comm_group + + self._compute_streams: List[torch.cuda.Stream] = self.create_compute_streams() + self._comm_stream: torch.cuda.Stream = torch.cuda.Stream(priority=-1) + + self._start_compute_event: torch.cuda.Event = torch.cuda.Event() + self._stop_compute_event: torch.cuda.Event = torch.cuda.Event() + + self._start_comm_event: torch.cuda.Event = torch.cuda.Event() + self._stop_comm_event: torch.cuda.Event = torch.cuda.Event() + + # keep origin state + self._main_stream: Optional[torch.cuda.Stream] = None + self._origin_ixf_comm_stream = None + self._ixformer_streams = dict() + + @classmethod + def dispatcher_key( + cls, num_chunks, num_compute_streams=None, comm_group=None, *args, **kwargs + ): + """ + the key of SplitOverlapComm + Args: + num_chunks: the number of chunks + num_compute_streams: the number of compute streams, default: 1 + comm_group: communicator group + Returns: unique key + """ + # warn: keey same function parameters with init + return (cls.__name__, num_chunks, num_compute_streams, comm_group) + + @classmethod + def enable(cls): + return config.IXFORMER_ENABLE_OVERLAP_COMM + + @property + def num_chunks(self): + return self._num_chunks + + @property + def num_compute_streams(self): + return self._num_compute_streams + + @property + def comm_group(self): + return self._comm_group + + def create_compute_streams(self): + streams = [] + for _ in range(self.num_compute_streams): + streams.append(torch.cuda.Stream()) + return streams + + def start_overlap(self): + self._main_stream = torch.cuda.current_stream() + + self._start_compute_event.record(torch.cuda.current_stream()) + for compute_stream in self._compute_streams: + compute_stream.wait_event(self._start_compute_event) + + self._origin_ixf_comm_stream = ixfd.get_comm_group_stream(self._comm_group) + ixfd.set_comm_group_stream(self._comm_stream.cuda_stream, self._comm_group) + + def stop_overlap(self): + last_compute_stream_id = ( + self.num_chunks + self.num_compute_streams - 1 + ) % self.num_compute_streams + self._stop_compute_event.record(self._compute_streams[last_compute_stream_id]) + self._stop_comm_event.record(self._comm_stream) + torch.cuda.current_stream().wait_event(self._stop_compute_event) + torch.cuda.current_stream().wait_event(self._stop_comm_event) + + ixfd.set_comm_group_stream(self._origin_ixf_comm_stream, self._comm_group) + + def start_comm(self, chunk_idx): + """ + prepare communication stream and wait event. + Args: + chunk_idx: the index of chunk + """ + + self._start_comm_event.record( + self._compute_streams[chunk_idx % self.num_compute_streams] + ) + self._comm_stream.wait_event(self._start_comm_event) + + @contextmanager + def compute_stream_context(self, chunk_idx): + """ + open python context and switch to compute stream in torch context + Args: + chunk_idx: the index of chunk + """ + + stream = self._compute_streams[chunk_idx % self.num_compute_streams] + + # print("before stream:", torch.cuda.current_stream()) + torch.cuda.set_stream(stream) + + # print("after stream:", torch.cuda.current_stream(), ixformer.cuda.current_stream()) + yield stream + + torch.cuda.set_stream(self._main_stream) + + @contextmanager + def stream_context(self, stream): + # print("before stream:", torch.cuda.current_stream()) + torch.cuda.set_stream(stream) + + # print("after stream:", torch.cuda.current_stream(), ixformer.cuda.current_stream()) + yield stream + + torch.cuda.set_stream(self._main_stream) + + def forward(self, *args, **kwargs): + self.start_overlap() + out = self.compute(*args, **kwargs) + self.stop_overlap() + return out + + @abc.abstractmethod + def compute(self, *args, **kwargs): + """ + it is abstract method to execute compute and communication. + """ + pass + + +class GemmMethod(enum.IntEnum): + kCUINFER = 0 + kCUBLAS = 1 + kLIMITED_GEMM = 2 + + +class GemmWithLimitedBlock: + def __init__(self, limit_algo=0) -> None: + self.limit_algo = limit_algo + self.env_key = "PYTORCH_GEMM_BLOCK_LIMITATION" + + def __enter__(self) -> None: + os.environ[self.env_key] = str(self.limit_algo) + + def __exit__(self, exc_type, exc_value, traceback) -> None: + del os.environ[self.env_key] + + +class IxFormerLimitedGemmContext: + def __init__(self) -> None: + self.env_key = "IXFORMER_ENABLE_PERSISTENT_GEMM" + + def __enter__(self) -> None: + os.environ[self.env_key] = "1" + + def __exit__(self, exc_type, exc_value, traceback) -> None: + os.environ[self.env_key] = "0" + + +class GemmAllReduceSplitOverlapComm(SplitOverlapComm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.gemm_method_env = config.IXFORMER_OVERLAP_GEMM_METHOD + + if self.gemm_method_env is None: + if ixfd.get_world_size(self.comm_group) == 2: + self.gemm_method_env = 0 + else: + self.gemm_method_env = 2 + + self.gemm_method = GemmMethod(int(self.gemm_method_env)) + self.limited_gemm_ctx = GemmWithLimitedBlock() + self.ixf_limited_gemm_ctx = IxFormerLimitedGemmContext() + self.split_ratio = config.IXFORMER_OVERLAP_SPLIT_RATIO + + @classmethod + def compute_row_parallel_dims(cls, input): + batch = 1 + if input.ndim == 2: + seqlen = input.shape[0] + else: + batch = input.shape[0] + seqlen = input.shape[1] + + parallel_dims = batch * seqlen + return parallel_dims + + def compute(self, input, weight, bias=None, out=None, *args, **kwargs): + """ + :param input: [Batch, SeqLen, Hidden] + :param weight: [OutChannel, InChannel] + :param bias: [OutChannel] + """ + + is_update_shape = input.ndim > 2 + batch = 1 + if input.ndim == 2: + seqlen = input.shape[0] + else: + batch = input.shape[0] + seqlen = input.shape[1] + + parallel_dims = batch * seqlen + + if is_update_shape: + input = input.reshape(parallel_dims, -1) + + if out is None: + out_shape = [parallel_dims, weight.shape[0]] + out_dtype = kwargs["out_dtype"] if "out_dtype" in kwargs else input.dtype + out = torch.empty(out_shape, dtype=out_dtype, device=input.device) + + if self.split_ratio is not None: + round_multiples = 256 if parallel_dims >= 256 else parallel_dims + first_chunk_size = ( + round((parallel_dims * float(self.split_ratio)) / round_multiples) + * round_multiples + ) + middle_chunk_size = (parallel_dims - first_chunk_size) // ( + self.num_chunks - 1 + ) + middle_chunk_size = (middle_chunk_size // round_multiples) * round_multiples + last_chunk_size = ( + parallel_dims + - first_chunk_size + - middle_chunk_size * (self.num_chunks - 2) + ) + + chunk_sizes = ( + [first_chunk_size] + + [middle_chunk_size] * (self.num_chunks - 2) + + [last_chunk_size] + ) + input_chunks = torch.split_with_sizes(input, chunk_sizes, dim=0) + out_chunks = torch.split_with_sizes(out, chunk_sizes, dim=0) + + # print(first_chunk_size, middle_chunk_size, last_chunk_size, chunk_sizes) + else: + input_chunks = torch.chunk(input, self.num_chunks, dim=0) + out_chunks = torch.chunk(out, self.num_chunks, dim=0) + + for chunk_idx in range(len(input_chunks)): + with self.compute_stream_context(chunk_idx): + chunk_out = self.gemm_dispatcher( + chunk_idx, + input_chunks[chunk_idx], + weight, + out_chunks[chunk_idx], + *args, + **kwargs, + ) + + self.start_comm(chunk_idx) + + ixfd.all_reduce( + chunk_out, async_op=True, group=self.comm_group, use_comm_stream=True + ) + + if is_update_shape: + out = out.reshape(batch, seqlen, -1) + + if bias is not None: + out = out + bias + + return out + + def gemm_dispatcher( + self, + chunk_idx, + chunk_input, + weight, + chunk_out=None, + user_gemm_method=None, + *args, + **kwargs, + ): + if user_gemm_method is not None and callable(user_gemm_method): + ctx = nullcontext() if chunk_idx == 0 else self.ixf_limited_gemm_ctx + with ctx: + return user_gemm_method( + chunk_input, weight, out=chunk_out, *args, **kwargs + ) + + if user_gemm_method is None: + user_gemm_method = self.gemm_method + + if user_gemm_method == GemmMethod.kCUINFER: + import ixformer.functions as ixff + + return ixff.linear(chunk_input, weight, output=chunk_out) + elif user_gemm_method == GemmMethod.kCUBLAS: + return torch.matmul(chunk_input, weight.T, out=chunk_out) + elif user_gemm_method == GemmMethod.kLIMITED_GEMM: + ctx = self.limited_gemm_ctx + with ctx: + return torch.matmul(chunk_input, weight.T, out=chunk_out) + elif user_gemm_method == GemmMethod.kCUBLAS: + return torch.matmul(chunk_input, weight.T, out=chunk_out) + else: + raise RuntimeError(f"Invalid gemm method, got {self.gemm_method}.") + + @classmethod + def native_forward( + cls, + input, + weight, + bias=None, + out=None, + group=None, + user_gemm_method=None, + *args, + **kwargs, + ): + if user_gemm_method is not None and callable(user_gemm_method): + gemm_out = user_gemm_method( + input, weight, bias=bias, out=out, *args, **kwargs + ) + out = out if gemm_out is None else gemm_out + else: + import ixformer.functions as ixff + + # warning: 下面的两种 gemm 可能存在精度不一致 + # out = torch.matmul(input, weight.T, out=out) + out = ixff.linear(input=input, weight=weight, bias=bias, output=out) + ixfd.all_reduce(out, async_op=True, group=group) + return out + + @classmethod + def is_supported(cls, input, num_chunks, comm_group): + if not cls.enable(): + return False + + ndim = input.ndim + shape = input.shape + + if ndim == 1: + m, k = 1, shape[0] + elif ndim == 2: + m, k = shape + else: + m, k = sum(shape[:-1]), shape[-1] + + return m >= 512 + + +_DEFAULT_OVERLAP_GROUP = None +_DEFAULT_OVERLAP_COMM_N2 = None +_DEFAULT_OVERLAP_COMM_N4 = None +_DEFAULT_OVERLAP_CHUNKS = config.IXFORMER_OVERLAP_CHUNKS + + +def linear_allreduce_overlap( + input, weight, bias=None, out=None, group=None, num_chunks=None, *args, **kwargs +): + num_chunks = num_chunks or _DEFAULT_OVERLAP_CHUNKS + + # print("call overlap:", GemmAllReduceSplitOverlapComm.is_supported(input, num_chunks=num_chunks, comm_group=group), input.shape, weight.shape if torch.is_tensor(weight) else None, "WorldSize:", ixfd.get_group_world_size(group), ", NumChunks:", num_chunks) + if not GemmAllReduceSplitOverlapComm.is_supported( + input, num_chunks=num_chunks, comm_group=group + ): + return GemmAllReduceSplitOverlapComm.native_forward( + input, weight, bias=bias, out=out, group=group, *args, **kwargs + ) + + global _DEFAULT_OVERLAP_GROUP + global _DEFAULT_OVERLAP_COMM_N2 + global _DEFAULT_OVERLAP_COMM_N4 + + if _DEFAULT_OVERLAP_GROUP is None: + _DEFAULT_OVERLAP_GROUP = group + + if num_chunks == 2 and group == _DEFAULT_OVERLAP_GROUP: + if _DEFAULT_OVERLAP_COMM_N2 is None: + _DEFAULT_OVERLAP_COMM_N2 = GemmAllReduceSplitOverlapComm.dispatcher( + num_chunks=num_chunks, comm_group=group + ) + overlap_comm = _DEFAULT_OVERLAP_COMM_N2 + elif num_chunks == 4 and group == _DEFAULT_OVERLAP_GROUP: + if _DEFAULT_OVERLAP_COMM_N4 is None: + _DEFAULT_OVERLAP_COMM_N4 = GemmAllReduceSplitOverlapComm.dispatcher( + num_chunks=num_chunks, comm_group=group + ) + overlap_comm = _DEFAULT_OVERLAP_COMM_N4 + else: + overlap_comm = GemmAllReduceSplitOverlapComm.dispatcher( + num_chunks=num_chunks, comm_group=group + ) + return overlap_comm.forward(input, weight, bias=bias, out=out, *args, **kwargs) diff --git a/ixformer_sdk/functions/__init__.py b/ixformer_sdk/functions/__init__.py new file mode 100644 index 00000000..96422ef9 --- /dev/null +++ b/ixformer_sdk/functions/__init__.py @@ -0,0 +1 @@ +from ..inference.functions import * diff --git a/ixformer_sdk/inference/__init__.py b/ixformer_sdk/inference/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/inference/distributed/__init__.py b/ixformer_sdk/inference/distributed/__init__.py new file mode 100644 index 00000000..796b8cc5 --- /dev/null +++ b/ixformer_sdk/inference/distributed/__init__.py @@ -0,0 +1 @@ +from .mpi_utils import * diff --git a/ixformer_sdk/inference/distributed/mpi_utils.py b/ixformer_sdk/inference/distributed/mpi_utils.py new file mode 100644 index 00000000..4b2431fc --- /dev/null +++ b/ixformer_sdk/inference/distributed/mpi_utils.py @@ -0,0 +1,21 @@ +import os + +from mpi4py import MPI + + +def get_world_size(comm=None): + if comm is None: + comm = MPI.COMM_WORLD + + return comm.Get_size() + + +def get_local_rank(comm=None): + return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]) + + +def get_rank(comm=None): + if comm is None: + comm = MPI.COMM_WORLD + + return comm.Get_rank() diff --git a/ixformer_sdk/inference/functions/__init__.py b/ixformer_sdk/inference/functions/__init__.py new file mode 100644 index 00000000..f471aed8 --- /dev/null +++ b/ixformer_sdk/inference/functions/__init__.py @@ -0,0 +1,44 @@ +from .act_and_mul import * +from .act_bias_mm import * +from .add import * +from .bert import * +from .bnb_dequant import * +from .bnb_double_quant import * +from .bnb_mm_dequant import * +from .bnb_qgemm import * +from .bnb_quant import * +from .bnb_rowcol_absmax import * +from .conv2d import * +from .cross_entropy_loss import * +from .flash_attn import * +from .flash_attn_lib import * +from .fused_rope import * +from .gemv import * +from .groupnorm import * +from .i8w8o32 import * +from .layernorm import * +from .lightllm import * +from .linalg import * +from .linear import * +from .lmdeploy import * +from .marlin import * +from .matmul import * +from .mla_fused import * +from .mm import * +from .moe import * +from .overlap_comm import * +from .paged_attention import * +from .quantized_linear import * +from .residual_bias import * +from .rms_norm import * +from .scaled_dot_product_attention import * +from .smoothquant import * +from .softmax import * +from .store_kv_cache import * +from .t5 import * +from .tgi import * +from .vllm import * +from .w8a8 import * +from .w8a16 import * +from .wi4a16 import * +from .wui4a16 import * diff --git a/ixformer_sdk/inference/functions/act_and_mul.py b/ixformer_sdk/inference/functions/act_and_mul.py new file mode 100644 index 00000000..50dc9be7 --- /dev/null +++ b/ixformer_sdk/inference/functions/act_and_mul.py @@ -0,0 +1,88 @@ +from typing import List, Union + +import ixformer._C as ops +import torch +import torch.nn.functional as NNF + +__all__ = ["ref_silu_and_mul", "ref_gelu_and_mul", "ref_gelu_tanh_and_mul", + "silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul"] + + +def ref_silu_and_mul(input: "torch.Tensor") -> torch.Tensor: + x1, x2 = input.chunk(chunks=2, dim=-1) + res = NNF.silu(x1) * x2 + return res + + +def ref_gelu_and_mul(input: "torch.Tensor", gate_first=True) -> torch.Tensor: + x1, x2 = input.chunk(chunks=2, dim=-1) + if gate_first: + res = NNF.gelu(x1) * x2 + else: + res = NNF.gelu(x2) * x1 + return res + + +def ref_gelu_tanh_and_mul(input: "torch.Tensor") -> torch.Tensor: + x1, x2 = input.chunk(chunks=2, dim=-1) + res = NNF.gelu(x1) * x2 + return res + + +def silu_and_mul(input: torch.Tensor, output: torch.Tensor = None): + + """ + Args: + input: (..., 2*hidden_size) torch.float16, torch.bfloat16, torch.float32 + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + + Returns: + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + """ + if output is None: + output_shape = list(input.shape) + output_shape[-1] = output_shape[-1] // 2 + output = input.new_empty(output_shape) + + ops.infer.silu_and_mul(input, output) + + return output + + +def gelu_and_mul(input: "torch.Tensor", output: torch.Tensor = None, gate_first=True): + + """ + Args: + input: (..., 2*hidden_size) torch.float16, torch.bfloat16, torch.float32 + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + gate_first: bool + Returns: + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + """ + if output is None: + output_shape = list(input.shape) + output_shape[-1] = output_shape[-1] // 2 + output = input.new_empty(output_shape) + + ops.infer.gelu_and_mul(input, output, gate_first) + + return output + + +def gelu_tanh_and_mul(input: torch.Tensor, output: torch.Tensor = None): + + """ + Args: + input: (..., 2*hidden_size) torch.float16, torch.bfloat16, torch.float32 + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + Returns: + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + """ + if output is None: + output_shape = list(input.shape) + output_shape[-1] = output_shape[-1] // 2 + output = input.new_empty(output_shape) + + ops.infer.gelu_tanh_and_mul(input, output) + + return output \ No newline at end of file diff --git a/ixformer_sdk/inference/functions/act_bias_mm.py b/ixformer_sdk/inference/functions/act_bias_mm.py new file mode 100644 index 00000000..3a8d4025 --- /dev/null +++ b/ixformer_sdk/inference/functions/act_bias_mm.py @@ -0,0 +1,89 @@ +import ixformer._C as ops +import torch +import torch.nn.functional as NNF + +__all__ = ["act_bias_mm", "ref_act_bias_mm"] + + +def ref_act_bias_mm( + mat1: torch.Tensor, + mat2: torch.Tensor, + bias: torch.Tensor = None, + scale: float = 1, + act_type: str = "none", + trans_format: str = "NN", +): + assert len(mat1.shape) >= 2 + assert len(mat2.shape) >= 2 + if trans_format == "NN": + if bias is not None: + output = torch.matmul(mat1, mat2) * scale + bias + else: + output = torch.matmul(mat1, mat2) * scale + else: + if bias is not None: + output = torch.matmul(mat1, mat2.transpose(-1, -2)) * scale + bias + else: + output = torch.matmul(mat1, mat2.transpose(-1, -2)) * scale + if act_type == "gelu": + output = NNF.gelu(output) + elif act_type == "relu": + output = NNF.relu(output) + elif act_type == "silu": + output = NNF.silu(output) + elif act_type == "none": + output = output + else: + raise NotImplementedError() + return output + + +def act_bias_mm( + mat1: torch.Tensor, + mat2: torch.Tensor, + bias: torch.Tensor = None, + output: torch.Tensor = None, + scale: float = 1, + act_type: str = "none", + trans_format: str = "NN", +): + """ + Args: + mat1: [m,k] or [batch_count,m,k] torch.float16 + mat2: [k,n] or [n,k] torch.float16 + 当trans_format为"NN"时[k,n], 当trans_format为"TN"时[n,k] + bias: [n] torch.float16 + output: [m,n] torch.float16 + scale: float + act_type: silu/gelu/relu/None str + 如果act_type不为None,则bias也不可以为None + trans_format: NN or TN str + Returns: + output: [m,n] torch.float16 + """ + assert len(mat1.shape) >= 2 + assert len(mat2.shape) >= 2 + if output is None: + output_shape = list(mat1.shape) + m = mat1.shape[-2] + if trans_format == "NN": + n = mat2.shape[-1] + else: + n = mat2.shape[-2] + output_shape[-2] = m + output_shape[-1] = n + output = mat1.new_empty(output_shape) + + add_bias = False + if bias is not None: + add_bias = True + + if add_bias: + ops.infer.act_bias_mm( + mat1, mat2, bias, output, add_bias, scale, act_type, trans_format + ) + else: + ops.infer.act_bias_mm( + mat1, mat2, mat1, output, add_bias, scale, act_type, trans_format + ) + return output diff --git a/ixformer_sdk/inference/functions/add.py b/ixformer_sdk/inference/functions/add.py new file mode 100644 index 00000000..38814773 --- /dev/null +++ b/ixformer_sdk/inference/functions/add.py @@ -0,0 +1,46 @@ +import ixformer._C as ops +import torch + +__all__ = [ + "ref_add", + "add", +] + + +def ref_add(input: torch.Tensor, other: torch.Tensor, out: torch.Tensor = None): + return torch.add(input, other, out=out) + + +def add(input: torch.Tensor, other: torch.Tensor, out: torch.Tensor = None): + """ + out = input + other + Support elementwise addition, but broadcasting is not supported yet. + Note: The dtype of input and other needs to be the same. + Args: + input: (...) torch.float32, torch.float16, torch.bfloat16 + other: (...) same as input + out: (...) same as input + Returns: + out: (...) same as input + """ + if input.dtype not in [torch.float16, torch.float32, torch.bfloat16]: + return torch.add(input, other, out=out) + if not input.is_contiguous() or not other.is_contiguous(): + return torch.add(input, other, out=out) + if out is not None and not out.is_contiguous(): + return torch.add(input, other, out=out) + + if input.dtype != other.dtype: + return torch.add(input, other, out=out) + if out is not None and out.dtype != input.dtype: + return torch.add(input, other, out=out) + + assert input.shape == other.shape, (f"broadcasting is not supported yet." + "input is {input.shape}, other is {other.shape}") + + if out is None: + out = torch.empty_like(input) + + ops.infer.add(input, other, out) + + return out diff --git a/ixformer_sdk/inference/functions/bert.py b/ixformer_sdk/inference/functions/bert.py new file mode 100644 index 00000000..68fd97c2 --- /dev/null +++ b/ixformer_sdk/inference/functions/bert.py @@ -0,0 +1,199 @@ +import ixformer._C as ops +import torch + +__all__ = [ + "ref_bert_embedding", + "bert_embedding", + "ref_bert_add_norm", + "bert_add_norm", + "ref_bert_unpack_start_end_logits", + "bert_unpack_start_end_logits", + "ref_bert_linear_residual", + "bert_linear_residual", +] + + +def ref_bert_embedding( + token_weight: torch.Tensor, + pos_weight: torch.Tensor, + type_weight: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + token_ids: torch.Tensor, + pos_ids: torch.Tensor, + type_ids: torch.Tensor, + epsilon: float = 1e-5, + out: torch.Tensor = None, +): + assert out is None + emd1 = torch.nn.functional.embedding(token_ids, token_weight) + emd2 = torch.nn.functional.embedding(pos_ids, pos_weight) + emd3 = torch.nn.functional.embedding(type_ids, type_weight) + + out = emd1 + emd2 + emd3 + out = torch.nn.functional.layer_norm(out, [out.shape[-1]], ln_weight, ln_bias) + return out + + +def bert_embedding( + token_weight: torch.Tensor, + pos_weight: torch.Tensor, + type_weight: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + token_ids: torch.Tensor, + pos_ids: torch.Tensor, + type_ids: torch.Tensor, + epsilon: float = 1e-5, + out: torch.Tensor = None, +): + """ + Args: + token_weight: (vocab_size, hidden_size) torch.float16, torch.bfloat16 + pos_weight: (pos_size, hidden_size) same as token_weight + type_weight: (type_size, hidden_size) same as token_weight + ln_weight: (hidden_size) same as token_weight + ln_bias: (hidden_size) same as token_weight + token_ids: (num_tokens) torch.int32, torch.int64 + pos_ids: (num_tokens) same as token_ids + type_ids: (num_tokens) same as token_ids + epsilon: float + out: (num_tokens, hidden_size) same as token_weight + Returns: + out: (num_tokens, hidden_size) same as token_weight + """ + if out is None: + out_shape = list(token_ids.shape) + hidden_size = token_weight.shape[-1] + out_shape.append(hidden_size) + out = token_weight.new_empty(out_shape) + + ops.infer.bert_embedding( + token_weight, + pos_weight, + type_weight, + ln_weight, + ln_bias, + token_ids, + pos_ids, + type_ids, + out, + epsilon, + ) + + return out + + +def ref_bert_add_norm( + input: torch.Tensor, + residual: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + epsilon: float = 1e-5, + out: torch.Tensor = None, +): + assert out is None + input = input + residual + return torch.nn.functional.layer_norm( + input, [input.shape[-1]], ln_weight, ln_bias, epsilon + ) + + +def bert_add_norm( + input: torch.Tensor, + residual: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + epsilon: float = 1e-5, + out: torch.Tensor = None, +): + """ + out = input + residual + out = add_norm(out, ln_weight, ln_bias, epsilon) + Args: + input: (num_tokens, hidden_size) torch.float16, torch.bfloat16 + residual: (num_tokens, hidden_size) same as input + ln_weight: (hidden_size) same as input + ln_bias: (hidden_size) same as input + epsilon: float + out: (num_tokens, hidden_size) same as input + Returns: + out: (num_tokens, hidden_size) same as input + """ + if out is None: + out = torch.empty_like(input) + ops.infer.bert_add_norm(input, residual, ln_weight, ln_bias, out, epsilon) + return out + + +def ref_bert_unpack_start_end_logits( + logits: torch.Tensor, + cu_seq_lens: torch.Tensor, + max_seq_len: int, + start_logits: torch.Tensor = None, + end_logits: torch.Tensor = None, +): + batch_size = cu_seq_lens.shape[0] - 1 + if start_logits is None: + start_logits = logits.new_empty([batch_size, max_seq_len]) + if end_logits is None: + end_logits = logits.new_empty([batch_size, max_seq_len]) + cu_seq_len_cpu = cu_seq_lens.detach().cpu() + for i in range(batch_size): + start_idx = cu_seq_len_cpu[i] + end_idx = cu_seq_len_cpu[i + 1] + cur_len = end_idx - start_idx + start_logits[i, :cur_len] = logits[start_idx:end_idx, 0] + end_logits[i, :cur_len] = logits[start_idx:end_idx, 1] + return start_logits, end_logits + + +def bert_unpack_start_end_logits( + logits: torch.Tensor, + cu_seq_lens: torch.Tensor, + max_seq_len: int, + start_logits: torch.Tensor = None, + end_logits: torch.Tensor = None, +): + """ + Args: + logits: (num_tokens, 2) torch.float16, torch.bfloat16 + cu_seq_lens: (batch_size+1) torch.int32, torch.int64 + max_seq_len: int + start_logits: (batch_size, max_seq_len) same as logits + end_logits: (batch_size, max_seq_len) same as logits + Returns: + start_logits: (batch_size, max_seq_len) same as logits + end_logits: (batch_size, max_seq_len) same as logits + """ + batch_size = cu_seq_lens.shape[0] - 1 + if start_logits is None: + start_logits = logits.new_empty([batch_size, max_seq_len]) + if end_logits is None: + end_logits = logits.new_empty([batch_size, max_seq_len]) + ops.infer.bert_unpack_start_end_logits( + logits, cu_seq_lens, start_logits, end_logits + ) + return start_logits, end_logits + + +def ref_bert_linear_residual( + input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, out: torch.Tensor +): + return torch.nn.functional.linear(input, weight, bias) + out + + +def bert_linear_residual( + input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, out: torch.Tensor +): + """ + Args: + input: (m, k) torch.float16, torch.bfloat16 + weight: (n, k) same as input + bias: (n) same as input + out: (m, n) same as input + Returns: + out: (m, n) same as input + """ + ops.infer.bert_linear_residual(input, weight, bias, out) + return out diff --git a/ixformer_sdk/inference/functions/bnb_dequant.py b/ixformer_sdk/inference/functions/bnb_dequant.py new file mode 100644 index 00000000..4bb12592 --- /dev/null +++ b/ixformer_sdk/inference/functions/bnb_dequant.py @@ -0,0 +1,55 @@ +from typing import List, Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = [ + "bnb_dequant", + "ref_bnb_dequant", +] + + +def ref_bnb_dequant( + qA: torch.Tensor, + SA: torch.Tensor, + training: bool = False, + scale: float = 127.0, + dequant_type: int = 0, +): + A = torch.empty(qA.shape, dtype = SA.dtype, device = SA.device) + if dequant_type == 0: + for i in range(qA.size(0)): + A[i:] = qA[i:] * (SA[i].to(torch.float) / scale).to(SA.dtype) + else: + for i in range(qA.size(1)): + A[:,i] = qA[:,i] * (SA[i].to(torch.float) / scale).to(SA.dtype) + + return A + + +def bnb_dequant( + qA: torch.Tensor, + SA: torch.Tensor, + training: bool = False, + scale: float = 127.0, + dequant_type: int = 0, +) -> torch.Tensor: + + """ + Args: + qA: (row, col) torch.int8 + dequant input + SA: (row) or (col) torch.half + scale vector + training: bool + scale: float + dequnt_type: int + 0 : every row shared a scale, SA shape : [row] + 1 : every col shared a scale, SA shape : [col] + Returns: + Tensor: (row, col) torch.half + dequant output + + """ + return ops.infer.bnb_dequant(qA, SA, scale, dequant_type) diff --git a/ixformer_sdk/inference/functions/bnb_double_quant.py b/ixformer_sdk/inference/functions/bnb_double_quant.py new file mode 100644 index 00000000..173abe7f --- /dev/null +++ b/ixformer_sdk/inference/functions/bnb_double_quant.py @@ -0,0 +1,175 @@ +from typing import List, Union + +import ixformer._C as ops +from torch.autograd.function import Function, FunctionCtx + +__all__ = ["bnb_double_quant"] + +import ctypes as ct + +import torch +from torch import Tensor + + +def get_ptr(A): + if A is None: + return None + else: + return ct.c_void_p(A.data.data_ptr()) + + +class COOSparseTensor: + def __init__(self, rows, cols, nnz, rowidx, colidx, values): + assert rowidx.dtype == torch.int + assert colidx.dtype == torch.int + assert values.dtype == torch.half + assert values.numel() == nnz + assert rowidx.numel() == nnz + assert colidx.numel() == nnz + + self.rows = rows + self.cols = cols + self.nnz = nnz + self.rowidx = rowidx + self.colidx = colidx + self.values = values + + +def coo_zeros(rows, cols, nnz, device, dtype=torch.half): + rowidx = torch.full(size=(nnz,), fill_value=0, dtype=torch.int, device=device) + + colidx = torch.full((nnz,), fill_value=0, dtype=torch.int, device=device) + values = torch.full((nnz,), fill_value=0, dtype=dtype, device=device) + return COOSparseTensor(rows, cols, nnz, rowidx, colidx, values) + + +def get_colrow_absmax( + A, row_stats=None, col_stats=None, nnz_block_ptr=None, threshold=0.0 +): + cols = A.shape[-1] + if len(A.shape) == 3: + rows = A.shape[0] * A.shape[1] + else: + rows = A.shape[0] + + col_tiles = (cols + 255) // 256 + tiled_rows = ((rows + 15) // 16) * 16 + if row_stats is None: + row_stats = torch.full( + size=(rows,), fill_value=-50000.0, dtype=torch.float, device=A.device + ) + if col_stats is None: + col_stats = torch.full( + size=(cols,), fill_value=-50000.0, dtype=torch.float, device=A.device + ) + + # if nnz_block_ptr is None and threshold > 0.0: + nnz_block_ptr = torch.full( + size=(tiled_rows * col_tiles + 1,), + fill_value=0, + dtype=torch.int, + device=A.device, + ) + + ops.infer.bnb_getColRowStats( + A, row_stats, col_stats, nnz_block_ptr, threshold, rows, cols + ) + + return row_stats, col_stats, nnz_block_ptr + + +# A : quant input shape : [row, col] shape:torch.half +def bnb_double_quant( + A: torch.Tensor, training: bool = False, threshold: float = 0.0 +) -> torch.Tensor: + + """ + Args: + A: (row, col) torch.float16 + quant input + training: bool + threshold: float + abs of element exceeds threshold will be ignored + Returns: + out_row: (row, col) torch.int8 + out_col: (row, col) torch.int8 + row_stats (row) torch.float + col_stats (col) torch.float + coo_tensor + """ + + assert A.dtype == torch.half + + cols = A.shape[-1] + if len(A.shape) == 3: + rows = A.shape[0] * A.shape[1] + else: + rows = A.shape[0] + + row_stats, col_stats, nnz_row_ptr = get_colrow_absmax(A, threshold=threshold) + + out_col = torch.full(size=A.shape, fill_value=0, dtype=torch.int8, device=A.device) + out_row = torch.full(size=A.shape, fill_value=0, dtype=torch.int8, device=A.device) + + coo_tensor = None + if threshold > 0.0: + nnz = nnz_row_ptr.cpu().numpy()[-1] + if nnz > 0: + coo_tensor = coo_zeros(A.shape[0], A.shape[1], nnz, A.device) + + ops.infer.bnb_doubleRowColQuant( + A, + row_stats, + col_stats, + out_col, + out_row, + coo_tensor.rowidx, + coo_tensor.colidx, + coo_tensor.values, + nnz_row_ptr, + threshold, + rows, + cols, + ) + val, idx = torch.sort(torch.Tensor(coo_tensor.rowidx.cpu().numpy())) + coo_tensor.rowidx = val + coo_tensor.colidx = torch.Tensor(coo_tensor.colidx.cpu().numpy())[idx].to( + torch.int32 + ) + coo_tensor.values = torch.Tensor(coo_tensor.values.cpu().numpy())[idx].to( + torch.half + ) + # coo_tensor.colidx = coo_tensor.colidx[idx] + # coo_tensor.values = coo_tensor.values[idx] + else: + ops.infer.bnb_doubleRowColQuant( + A, + row_stats, + col_stats, + out_col, + out_row, + out_row, + out_row, + out_row, + out_row, + 0.0, + rows, + cols, + ) + else: + ops.infer.bnb_doubleRowColQuant( + A, + row_stats, + col_stats, + out_col, + out_row, + out_row, + out_row, + out_row, + out_row, + threshold, + rows, + cols, + ) + + return out_row, out_col, row_stats, col_stats, coo_tensor diff --git a/ixformer_sdk/inference/functions/bnb_mm_dequant.py b/ixformer_sdk/inference/functions/bnb_mm_dequant.py new file mode 100644 index 00000000..9fa02055 --- /dev/null +++ b/ixformer_sdk/inference/functions/bnb_mm_dequant.py @@ -0,0 +1,73 @@ +from typing import List, Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = ["bnb_mm_dequant"] + + +# A : quant input shape : [row, col] shape : torch.int +def bnb_mm_dequant( + A: torch.Tensor, + quant_state: tuple, + row_stats: torch.Tensor, + col_stats: torch.Tensor, + bias: torch.Tensor = None, + add_bias: bool = False, + training: bool = False, +) -> torch.Tensor: + + """ + Args: + A: (row, col) torch.int8 + quant_state: tuple + row_stats: (row) torch.float + col_stats: (col) torch.float + bias: (col) torch.half + add_bias: bool + training: bool + Returns: + Tensor: (row, col) torch.half + + """ + + assert A.dtype == torch.int + if bias is not None: + add_bias = True + print("bias.dtype:", bias.dtype) + assert bias.dtype == torch.half + else: + bias = A + out_shape = quant_state[0] + if len(out_shape) == 3: + out_shape = (out_shape[0] * out_shape[1], out_shape[2]) + out = torch.full(size=out_shape, fill_value=0, dtype=torch.half, device=A.device) + new_row_stats = torch.full( + size=(out_shape[0],), fill_value=0, dtype=torch.float, device=A.device + ) + new_col_stats = torch.full( + size=(out_shape[1],), fill_value=0, dtype=torch.float, device=A.device + ) + + assert ( + new_row_stats.shape[0] == row_stats.shape[0] + ), f"{new_row_stats.shape} vs {row_stats.shape}" + assert ( + new_col_stats.shape[0] == col_stats.shape[0] + ), f"{new_col_stats.shape} vs {col_stats.shape}" + numRows = out_shape[0] + numCols = out_shape[1] + ops.infer.bnb_mm_dequant( + A, + row_stats, + col_stats, + out, + new_row_stats, + new_col_stats, + numRows, + numCols, + add_bias, + bias, + ) + return out diff --git a/ixformer_sdk/inference/functions/bnb_qgemm.py b/ixformer_sdk/inference/functions/bnb_qgemm.py new file mode 100644 index 00000000..d9aa39b7 --- /dev/null +++ b/ixformer_sdk/inference/functions/bnb_qgemm.py @@ -0,0 +1,56 @@ +from typing import List, Union + +import ixformer._C as ops +import torch + +__all__ = ["bnb_qgemm", "ref_bnb_qgemm"] + + +# qA : quant input shape : [bs, in_feature] +# qW : quant weight shape : [out_feature, in_feature] +# SA : scale vector of qA shape : [bs] +# SW : scale vector of qW shape : [out_feature] + + +def ref_bnb_qgemm( + qA: torch.Tensor, + qW: torch.Tensor, + SA: torch.Tensor, + SW: torch.Tensor, + training: bool = False, + scaleA: float = 127.0, + scaleW: float = 127.0, +): + y = torch.nn.functional.linear(qA.to(torch.float), qW.to(torch.float)) + out = torch.empty(y.shape, dtype = SA.dtype, device = SA.device) + for i in range(qA.size(0)): + for j in range(qW.size(0)): + out[i][j] = y[i][j] * (SA[i].to(torch.float) / scaleA) * (SW[j].to(torch.float) / scaleW) + return out.to(SA.dtype) + + +def bnb_qgemm( + qA: torch.Tensor, + qW: torch.Tensor, + SA: torch.Tensor, + SW: torch.Tensor, + training: bool = False, + scaleA: float = 127.0, + scaleW: float = 127.0, +) -> torch.Tensor: + + """ + Args: + qA: (bs, in_feature) torch.int8 + qW: (out_feature, in_feature) torch.int8 + SA: (bs) torch.half + scale vector of qA + SA: (out_feature) torch.half + scale vector of qW + training: bool + scaleA: float + scaleW: float + Returns: + Tensor: (bs, out_feature) torch.half + """ + return ops.infer.bnb_qgemm(qA, qW, SA, SW, scaleA, scaleW) diff --git a/ixformer_sdk/inference/functions/bnb_quant.py b/ixformer_sdk/inference/functions/bnb_quant.py new file mode 100644 index 00000000..5aa6014b --- /dev/null +++ b/ixformer_sdk/inference/functions/bnb_quant.py @@ -0,0 +1,58 @@ +from typing import List, Union + +import ixformer._C as ops +import torch + +__all__ = ["bnb_quant", "ref_bnb_quant"] + + +# A : input shape : [row, col] +# SA : scale vector +# quant_type +# 0 : every row shared a scale, SA shape : [row] +# 1 : every col shared a scale, SA shape : [col] +def ref_bnb_quant( + A: torch.Tensor, + SA: torch.Tensor, + training: bool = False, + scale: float = 127.0, + quant_type: int = 0, +): + qA = torch.empty(A.shape, device = SA.device) + if quant_type == 0: + for i in range(A.size(0)): + qA[i:] = torch.round(A[i:] * (scale / SA[i].to(torch.float))) + else: + for i in range(A.size(1)): + qA[:,i] = torch.round(A[:,i] * (scale / SA[i].to(torch.float))) + + qA_clamped = torch.clamp(qA, min=-128, max=127) + qA = qA_clamped.to(torch.int8) + return qA + + +def bnb_quant( + A: torch.Tensor, + SA: torch.Tensor, + training: bool = False, + scale: float = 127.0, + quant_type: int = 0, +) -> torch.Tensor: + + """ + Args: + A: (row, col) torch.half + quant input + SA: (row) or (col) torch.half + scale vector + training: bool + scale: float + qunt_type: int + 0 : every row shared a scale, SA shape : [row] + 1 : every col shared a scale, SA shape : [col] + Returns: + Tensor: (row, col) torch.int8 + quant output + + """ + return ops.infer.bnb_quant(A, SA, scale, quant_type) diff --git a/ixformer_sdk/inference/functions/bnb_rowcol_absmax.py b/ixformer_sdk/inference/functions/bnb_rowcol_absmax.py new file mode 100644 index 00000000..aeeda80f --- /dev/null +++ b/ixformer_sdk/inference/functions/bnb_rowcol_absmax.py @@ -0,0 +1,52 @@ +from typing import List, Union + +import ixformer._C as ops +import torch + +__all__ = ["bnb_rowcol_absmax", "ref_bnb_rowcol_absmax"] + + +# input : input shape : [row, col] +# threshold : abs of element exceeds threshold will be ignored +# type +# 0 : row absmax +def ref_bnb_rowcol_absmax( + input: torch.Tensor, + training: bool = False, + threshold: float = 0.0, + type: int = 0, +): + input = input.float() + if threshold ==0.0: + threshold = float('inf') + mask = (torch.abs(input) < threshold) + masked_input = mask * input + masked_input = masked_input.half() + if type == 0: + out = torch.amax(torch.abs(masked_input), dim=1) + + else: + out = torch.amax(torch.abs(masked_input), dim=0) + return out + + +def bnb_rowcol_absmax( + input: torch.Tensor, + training: bool = False, + threshold: float = 0.0, + type: int = 0, +) -> torch.Tensor: + + """ + Args: + input: (row, col) torch.half + 目前col值必须满足col%2==0 + training: bool + threshold: float + abs of element exceeds threshold will be ignored + type: int + row absmax, 目前只支持type=0 + Returns: + Tensor: (row) torch.half + """ + return ops.infer.bnb_rowcol_absmax(input, threshold, type) diff --git a/ixformer_sdk/inference/functions/conv2d.py b/ixformer_sdk/inference/functions/conv2d.py new file mode 100644 index 00000000..9b02fa5c --- /dev/null +++ b/ixformer_sdk/inference/functions/conv2d.py @@ -0,0 +1,198 @@ +from typing import Union + +import ixformer._C as ops +import torch +import torch.nn.functional as NNF + +__all__ = ["conv2d", "ref_conv2d", "ref_conv2d_nhwc", "conv2d_nhwc"] + + +def is_channels_last(ten): + return torch._prims_common.suggest_memory_format(ten) == torch.channels_last + + +def _pair(x): + if isinstance(x, (list, tuple)): + return x + return (x, x) + + +def ref_conv2d( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + stride: Union[int, tuple] = 1, + padding: Union[int, tuple] = 0, + dilation: Union[int, tuple] = 1, + groups: int = 1, +): + + output = NNF.conv2d(input, weight, bias, stride, padding, dilation, groups) + return output + + +# conv2d官方接口，如果weight是torch.channels_last,输出也是torch.channels_last；如果weight是nchw，那么输出也是nchw;特殊情况，如果输入是nchw，weight是torch.channels_last，输出也是torch.channels_last +def conv2d( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + stride: Union[int, tuple] = 1, + padding: Union[int, tuple] = 0, + dilation: Union[int, tuple] = 1, + groups: int = 1, +): + """ + Args: + input: (n,in_c,h,w) torch.float16 + weight: (out_c,in_c/groups,kH,kW) torch.float16 + bias: (out_c) torch.float16 + stride: int or tuple + Stride of the convolution. Default: 1 + padding: int or tuple + Padding added to all four sides of the input. Default: 0 + dilation: int or tuple + Spacing between kernel elements. Default: 1 + groups: int + Number of blocked connections from input channels to output channels. Default: 1 + Returns: + Tensor: (n,out_c,h_out,w_out) torch.float16 + h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) / stride_h + 1; + w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1; + """ + stride = _pair(stride) + padding = _pair(padding) + dilation = _pair(dilation) + + channel_last = is_channels_last(weight) + if not is_channels_last(input) and channel_last: + input = input.to(memory_format=torch.channels_last) + + # compute outshape + n, in_c, h_in, w_in = input.shape + out_c, _, kernel_h, kernel_w = weight.shape + pad_h = padding[0] + pad_w = padding[1] + stride_h = stride[0] + stride_w = stride[1] + dilation_h = dilation[0] + dilation_w = dilation[1] + h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) // stride_h + 1 + w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) // stride_w + 1 + + if channel_last: + output_shape = [n, out_c, h_out, w_out] + output = torch.empty( + output_shape, + memory_format=torch.channels_last, + dtype=input.dtype, + device=input.device, + ) + else: + output_shape = [n, out_c, h_out, w_out] + output = input.new_empty(output_shape) + + if channel_last: + input = input.permute(0, 2, 3, 1) + weight = weight.permute(0, 2, 3, 1) + output = output.permute(0, 2, 3, 1) + + if bias is not None: + bias = bias.float() + ops.infer.conv2d( + input, weight, bias, output, stride, padding, dilation, groups, channel_last + ) + if channel_last: + output = output.permute(0, 3, 1, 2) + + return output + + +def ref_conv2d_nhwc( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + stride: Union[int, tuple] = 1, + padding: Union[int, tuple] = 0, + dilation: Union[int, tuple] = 1, + groups: int = 1, +): + + output = NNF.conv2d( + input.permute(0, 3, 1, 2).contiguous(), + weight.permute(0, 3, 1, 2).contiguous(), + bias, + stride, + padding, + dilation, + groups, + ) + return output.permute(0, 2, 3, 1).contiguous() + + +# conv2d_nhwc， +# conv2d官方接口解决两种情况： +# 1、务必输入tensor内存上是nhwc,且tensor属于memory_format=torch.channels_last， +# 2、或者输入tensor内存上是nchw，并且是contiguous； +# conv2d官方接口不能解决，conv2d_nhwc则可处理这种情况的 +# 输入tensor内存上是nhwc的，但tensor没有用memory_format=torch.channels_last进行过处理，不会有memory_format=torch.channels_last的标签 +def conv2d_nhwc( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + stride: Union[int, tuple] = 1, + padding: Union[int, tuple] = 0, + dilation: Union[int, tuple] = 1, + groups: int = 1, +): + + """ + Args: + input: (n,h,w,in_c) torch.float16 + weight: (out_c,kH,kW,in_c/groups) torch.float16 + bias: (out_c) torch.float16 + stride: int or tuple + Stride of the convolution. Default: 1 + padding: int or tuple + Padding added to all four sides of the input. Default: 0 + dilation: int or tuple + Spacing between kernel elements. Default: 1 + groups: int + Number of blocked connections from input channels to output channels. Default: 1 + Returns: + Tensor: (n,h_out,w_out,out_c) torch.float16 + h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) / stride_h + 1; + w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1; + """ + + stride = _pair(stride) + padding = _pair(padding) + dilation = _pair(dilation) + + assert input.is_contiguous() + assert weight.is_contiguous() + + # compute outshape + n, h_in, w_in, in_c = input.shape + ( + out_c, + kernel_h, + kernel_w, + _, + ) = weight.shape + pad_h = padding[0] + pad_w = padding[1] + stride_h = stride[0] + stride_w = stride[1] + dilation_h = dilation[0] + dilation_w = dilation[1] + h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) // stride_h + 1 + w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) // stride_w + 1 + + output_shape = [n, h_out, w_out, out_c] + output = torch.empty(output_shape, dtype=input.dtype, device=input.device) + if bias is not None: + bias = bias.float() + ops.infer.conv2d( + input, weight, bias, output, stride, padding, dilation, groups, True + ) + return output diff --git a/ixformer_sdk/inference/functions/cross_entropy_loss.py b/ixformer_sdk/inference/functions/cross_entropy_loss.py new file mode 100644 index 00000000..359d6b1c --- /dev/null +++ b/ixformer_sdk/inference/functions/cross_entropy_loss.py @@ -0,0 +1,204 @@ +from typing import Union + +import ixformer._C as ops +import torch + +__all__ = ["vocab_parallel_cross_entropy", "ref_vocab_parallel_cross_entropy"] + + +def ref_vocab_parallel_cross_entropy( + vocab_parallel_logits: torch.Tensor, + target: torch.Tensor, + label_smoothing: float = 0.0, + world_size: int = 1, + vocab_start_index: int = 0, + vocab_end_index: int = 320000, + group=None, +): + if world_size == 1: + vocab_parallel_logits = vocab_parallel_logits.float() + partition_vocab_size = vocab_parallel_logits.size()[-1] + logits_max = torch.max(vocab_parallel_logits, dim=-1)[0] + vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1) + target_mask = (target < vocab_start_index) | (target >= vocab_end_index) + masked_target = target.clone() - vocab_start_index + masked_target[target_mask] = 0 + logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size) + masked_target_1d = masked_target.view(-1) + arange_1d = torch.arange( + start=0, end=logits_2d.size()[0], device=logits_2d.device + ) + + predicted_logits_1d = logits_2d[arange_1d, masked_target_1d] + predicted_logits_1d = predicted_logits_1d.clone().contiguous() + + predicted_logits = predicted_logits_1d.view_as(target) + predicted_logits[target_mask] = 0.0 + exp_logits = vocab_parallel_logits + torch.exp(vocab_parallel_logits, out=exp_logits) + sum_exp_logits = exp_logits.sum(dim=-1) + loss = torch.log(sum_exp_logits) - predicted_logits + exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1)) + if label_smoothing > 0: + """ + We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth. + = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt}) + = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i + = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K + From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py + """ + assert 1.0 > label_smoothing > 0.0 + smoothing = label_smoothing * partition_vocab_size / (partition_vocab_size - 1) + + # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs. + log_probs = torch.log(exp_logits) + mean_log_probs = log_probs.mean(dim=-1) + loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs + + else: + # Maximum value along vocab dimension across all GPUs. + logits_max = torch.max(vocab_parallel_logits, dim=-1)[0] + torch.distributed.all_reduce( + logits_max, op=torch.distributed.ReduceOp.MAX, group=group + ) + # Subtract the maximum value. + vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1) + + # Get the partition's vocab indecies + partition_vocab_size = vocab_parallel_logits.size()[-1] + + # Create a mask of valid vocab ids (1 means it needs to be masked). + target_mask = (target < vocab_start_index) | (target >= vocab_end_index) + masked_target = target.clone() - vocab_start_index + masked_target[target_mask] = 0 + + # Get predicted-logits = logits[target]. + # For Simplicity, we convert logits to a 2-D tensor with size + # [*, partition-vocab-size] and target to a 1-D tensor of size [*]. + logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size) + masked_target_1d = masked_target.view(-1) + arange_1d = torch.arange( + start=0, end=logits_2d.size()[0], device=logits_2d.device + ) + predicted_logits_1d = logits_2d[arange_1d, masked_target_1d] + predicted_logits_1d = predicted_logits_1d.clone().contiguous() + predicted_logits = predicted_logits_1d.view_as(target) + predicted_logits[target_mask] = 0.0 + # All reduce is needed to get the chunks from other GPUs. + torch.distributed.all_reduce( + predicted_logits, + op=torch.distributed.ReduceOp.SUM, + group=group, + ) + + # Sum of exponential of logits along vocab dimension across all GPUs. + exp_logits = vocab_parallel_logits + torch.exp(vocab_parallel_logits, out=exp_logits) + sum_exp_logits = exp_logits.sum(dim=-1) + torch.distributed.all_reduce( + sum_exp_logits, + op=torch.distributed.ReduceOp.SUM, + group=group, + ) + + # Loss = log(sum(exp(logits))) - predicted-logit. + loss = torch.log(sum_exp_logits) - predicted_logits + + # Normalize and optionally smooth logits + exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1)) + + vocab_size = exp_logits.size(-1) + if label_smoothing > 0: + """ + We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth. + = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt}) + = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i + = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K + From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py + """ + assert 1.0 > label_smoothing > 0.0 + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + + # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs. + log_probs = torch.log(exp_logits) + mean_log_probs = log_probs.mean(dim=-1) + loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs + + return loss + + +def vocab_parallel_cross_entropy( + vocab_parallel_logits: torch.Tensor, + target: torch.Tensor, + label_smoothing: float = 0.0, + world_size: int = 1, + vocab_start_index: int = 0, + vocab_end_index: int = 320000, + group=None, +): + """ + Args: + vocab_parallel_logits: (seq_len,1,vocal_size) torch.float16, torch.bfloat16, torch.float + target: (seq_len,1) torch.int64 + label_smoothing: float + 默认为0.0. 用于标签平滑 + world_size: int + 当world_size = 1时,目前只支持batch_size = 1 的情况 + vocab_start_index: int + vocab_end_index: int + group: + TP 并行组 + Returns: + loss: (seq_len,1) torch.float + """ + if world_size == 1: + device = vocab_parallel_logits.device + xnumel = vocab_parallel_logits.shape[0] + rnumel = vocab_parallel_logits.shape[-1] + exp_logits = torch.empty( + (xnumel, 1, rnumel), device=device, dtype=torch.float32 + ) + masked_target_1d = torch.empty((xnumel,), device=device, dtype=torch.int32) + loss = torch.empty((xnumel, 1), device=device, dtype=torch.float32) + + ops.train.cross_entropy_loss_forward( + vocab_parallel_logits, target.int(), exp_logits, masked_target_1d, loss + ) + + vocab_size = exp_logits.size(-1) + if label_smoothing > 0: + """ + We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth. + = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt}) + = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i + = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K + From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py + """ + assert 1.0 > label_smoothing > 0.0 + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + + # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs. + log_probs = torch.log(exp_logits) + mean_log_probs = log_probs.mean(dim=-1) + loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs + + # Store softmax, target-mask and masked-target for backward pass. + + return loss + else: + loss = ref_vocab_parallel_cross_entropy( + vocab_parallel_logits, + target, + label_smoothing, + world_size, + vocab_start_index, + vocab_end_index, + group, + ) + return loss diff --git a/ixformer_sdk/inference/functions/flash_attn.py b/ixformer_sdk/inference/functions/flash_attn.py new file mode 100644 index 00000000..9b360036 --- /dev/null +++ b/ixformer_sdk/inference/functions/flash_attn.py @@ -0,0 +1,201 @@ +import math +from typing import List, Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = [ + "ixinfer_flash_attn_unpad", + "ixinfer_flash_attn_pad", + "ref_ixinfer_flash_attn_pad", +] + + +def ixinfer_flash_attn_unpad( + # total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i + q: "torch.Tensor", + # total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + k: "torch.Tensor", + # total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i + v: "torch.Tensor", + # total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i + cu_seqlens_q: "torch.Tensor", # b+1 + cu_seqlens_k: "torch.Tensor", # b+1 + max_seqlen_q: int, + max_seqlen_k: int, + is_causal: bool = False, + atten_scale: float = None, + sqrt_alibi: bool = False, + # total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i + alibi_slopes: "torch.Tensor" = None, + out: "torch.Tenosr" = None, +): + """ + Args: + q: (total_q, nheads, headdim) torch.float16, torch.bfloat16 + where total_q = total number of query tokens in the batch. + k: (total_k, nheads_k, headdim) torch.float16, torch.bfloat16 + where total_k = total number of key tokens in the batch. + v: (total_k, nheads_k, headdim) torch.float16, torch.bfloat16 + cu_seqlens_q: (batch_size + 1) torch.int32 + The cumulative sequence lengths of the sequences in the batch, used to index into q. + cu_seqlens_k: (batch_size + 1) torch.int32 + The cumulative sequence lengths of the sequences in the batch, used to index into kv. + max_seqlen_q: int + Maximum query sequence length in the batch. + max_seqlen_k: int + Maximum key sequence length in the batch. + atten_scale: float + The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim). + is_causal: bool + Whether to apply causal attention mask (e.g., for auto-regressive modeling). + sqrt_alibi: bool + Whether to apply abilimode + out: (total, nheads, headdim) torch.float16, torch.bfloat16 + Returns: + out: (total, nheads, headdim) torch.float16, torch.bfloat16 + if not q.size(-1) % 32 == 0: out shape is (total_q, nheads, q.size(-1) + (32 - q.size(-1) % 32)) + """ + if atten_scale is None: + atten_scale = 1.0 / (q.size(-1) ** 0.5) + + # 判断是否pad + cur_head = q.size(-1) + cur_head32 = cur_head + if not cur_head % 32 == 0: + cur_head32 = cur_head + (32 - cur_head % 32) + q_infer = torch.nn.functional.pad(q, [0, cur_head32 - cur_head, 0, 0], value=0) + k_infer = torch.nn.functional.pad(k, [0, cur_head32 - cur_head, 0, 0], value=0) + v_infer = torch.nn.functional.pad(v, [0, cur_head32 - cur_head, 0, 0], value=0) + else: + q_infer = q + k_infer = k + v_infer = v + + if out is None: + out = torch.empty_like(q_infer) + # ixinfer 新接口版 + ops.infer.ixinfer_flash_attn_unpad( + q_infer, + k_infer, + v_infer, + out, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + is_causal, + False, # need_lse =False + atten_scale, + sqrt_alibi, + alibi_slopes, + ) + if not cur_head % 32 == 0: + out = out[:, :, :cur_head] + return out + + +def ref_ixinfer_flash_attn_pad( + # [ batch num_heads seq_q head_size] + q: torch.Tensor, + # [ batch num_heads_k max_seq_kv head_size] + k: torch.Tensor, + # [ batch num_heads_k max_seq_kv head_size] + v: torch.Tensor, + # [ batch num_heads seq_q seq_kv] seq_kv<=max_seq_kv + mask: torch.Tensor, + # [ batch num_heads seq_q head_size] + atten_scale: float = None, + kv_seq_start: int = None, + kv_seq_end: int = None, +): + head_dim = q.size(-1) + k_effective = k[:, :, kv_seq_start:kv_seq_end, :] + v_effective = v[:, :, kv_seq_start:kv_seq_end, :] + # 2. q*kt softmax + scores_qk = ( + torch.matmul(q.float(), k_effective.float().transpose(-2, -1)) * atten_scale + ) + # softmax + # print(scores_qk.shape,mask.shape) + if mask is not None: + if mask.dtype == torch.int32: + scores_qk = scores_qk + mask * (-100000) + elif mask.dtype == torch.float32: + scores_qk = scores_qk + mask + else: + print( + f"mask dtype is not surported {mask.dtype},now surport int32 and float32" + ) + scores_qk = torch.nn.functional.softmax(scores_qk, dim=-1) + # 3. x = qk_scores * v + scores_v = torch.matmul(scores_qk, v_effective.float()) + return scores_v.half() + + +def ixinfer_flash_attn_pad( + # [ batch num_heads seq_q head_size] + q: torch.Tensor, + # [ batch num_heads_k max_seq_kv head_size] + k: torch.Tensor, + # [ batch num_heads_k max_seq_kv head_size] + v: torch.Tensor, + # [ batch num_heads seq_q seq_kv] seq_kv<=max_seq_kv + mask: torch.Tensor, + # [ batch num_heads seq_q head_size] + atten_scale: float = None, + kv_seq_start: int = None, + kv_seq_end: int = None, +): + """ + Args: + q: (batch_size, num_head, seq_len_q, head_dim) torch.float16, torch.bfloat16 + k: (batch_size, num_head_kv, seq_len_kv, head_dim) torch.float16, torch.bfloat16 + v: (batch_size, num_head_kv, seq_len_kv, head_dim) torch.float16, torch.bfloat16 + mask: (batch_size, num_head, seq_len_q, kv_seq_start:kv_seq_end) torch.int32, torch.int64, torch.float32 + atten_scale: float + The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim). + kv_seq_start: int + kv sequence start index used for computation in the batch + kv_seq_end: int + kv sequence end index used for computation in the batch. + Returns: + out: (batch_size, num_head, seq_len_q, head_dim) torch.float16, torch.bfloat16 + """ + + # 判断是否pad + cur_head = q.size(-1) + cur_head32 = cur_head + if not cur_head % 32 == 0: + cur_head32 = cur_head + (32 - cur_head % 32) + q_infer = torch.nn.functional.pad(q, [0, cur_head32 - cur_head, 0, 0], value=0) + k_infer = torch.nn.functional.pad(k, [0, cur_head32 - cur_head, 0, 0], value=0) + v_infer = torch.nn.functional.pad(v, [0, cur_head32 - cur_head, 0, 0], value=0) + else: + q_infer = q + k_infer = k + v_infer = v + + if atten_scale is None: + atten_scale = 1.0 / (q.size(-1) ** 0.5) + if kv_seq_start is None or kv_seq_end is None: + kv_seq_start = 0 + kv_seq_end = k.size(-2) # kv seq len + elif kv_seq_start < 0 or kv_seq_end > k.size(-2) or kv_seq_start >= kv_seq_end: + raise NotImplementedError( + "must kv_seq_start<0 or kv_seq_end>k.size(-2) or kv_seq_start>=kv_seq_end!" + ) + out_shape = list(q_infer.shape) + out = torch.empty(out_shape, dtype=q.dtype, device=q.device) + if mask is not None: + ops.infer.ixinfer_flash_attn_pad_fwd( + q_infer, k_infer, v_infer, mask, out, atten_scale, kv_seq_start, kv_seq_end + ) + else: + ops.infer.ixinfer_flash_attn_pad_fwd_nomask( + q_infer, k_infer, v_infer, out, atten_scale, kv_seq_start, kv_seq_end + ) + if not cur_head % 32 == 0: + out = out[:, :, :, :cur_head] + return out diff --git a/ixformer_sdk/inference/functions/flash_attn_lib.py b/ixformer_sdk/inference/functions/flash_attn_lib.py new file mode 100644 index 00000000..a3966fa6 --- /dev/null +++ b/ixformer_sdk/inference/functions/flash_attn_lib.py @@ -0,0 +1,350 @@ +import math +from typing import List, Union + +import ixformer._C as ops +import torch + +from .flash_attn import ixinfer_flash_attn_unpad + +__all__ = [ + "flash_attn_varlen_func", + "ref_flash_attn_varlen_func", + "flash_attn_func", + "ref_flash_attn_func", +] + + +def ref_flash_attn_varlen_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + dropout_p: float = 0.0, + softmax_scale: float = None, + causal: bool = False, + return_attn_probs: bool = False, +): + if return_attn_probs: + raise NotImplementedError("return_attn_probs not supported!") + out = torch.zeros_like(q) + unpad_causal_torch( + q, + k, + v, + out, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + torch.float16, + softmax_scale, + causal, + ) + return out + + +def unpad_causal_torch( + q, + k, + v, + output, + cu_seqlens_q, + cu_seqlens_k, + max_seq_len_q, + max_seq_len_kv, + dtype, + atten_scale, + is_causal=True, +): + + head_num = q.size(1) + head_num_kv = k.size(1) + head_dim = q.size(2) + + assert head_num % head_num_kv == 0 + if atten_scale == None: + atten_scale = 1.0 / (q.size(-1) ** 0.5) + # tokens,head_num,head_dim + if head_num != head_num_kv: + # k = k.repeat(1, head_num//head_num_kv, 1)#[0,1,2,0,1,2,0,1,2,0,1,2] + # v = v.repeat(1, head_num//head_num_kv, 1) + + k = repeat_kv(k, head_num // head_num_kv) # [0,0,0,0,1,1,1,1,2,2,2,2] GROUP + v = repeat_kv(v, head_num // head_num_kv) + + batch_size = cu_seqlens_q.size(0) - 1 + + for i in range(batch_size): + q_start_index = cu_seqlens_q[i] + q_end_index = cu_seqlens_q[i + 1] + cur_q_len = q_end_index - q_start_index + # 1*seq_len,head_num,head_dim + cur_q = q[q_start_index:q_end_index] + + k_start_index = cu_seqlens_k[i] + k_end_index = cu_seqlens_k[i + 1] + cur_k_len = k_end_index - k_start_index + + cur_k = k[k_start_index:k_end_index] + cur_v = v[k_start_index:k_end_index] + + # mask = torch.tril(torch.ones([cur_q_len, cur_k_len], dtype=torch.bool)).cuda() + # mask = mask.unsqueeze(0).unsqueeze(0) + if is_causal: + # Create attention mask. + attn_mask = torch.triu( + torch.ones(cur_q_len, cur_k_len, dtype=dtype), diagonal=1 + ) + attn_mask = attn_mask * torch.finfo(dtype).min + attn_mask = attn_mask.to(dtype=dtype, device="cuda") + else: + attn_mask = None + + ref_output = ref_masked_attention( + cur_q, + cur_k, + cur_v, + atten_scale, + attn_mask=attn_mask, + ) + output[q_start_index:q_end_index].copy_(ref_output) + + +def ref_masked_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + attn_mask=None, +) -> torch.Tensor: + query = query * scale + dtype = query.dtype + device = query.device + query = query.to(torch.float32).cpu() + key = key.to(torch.float32).cpu() + value = value.to(torch.float32).cpu() + attn = torch.einsum("qhd,khd->hqk", query, key) + if attn_mask is not None: + attn_mask = attn_mask.cpu() + attn = attn + attn_mask + attn = torch.softmax(attn, dim=-1) + out = torch.einsum("hqk,khd->qhd", attn, value) + out = out.to(device).to(dtype) + return out + + +def flash_attn_varlen_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + dropout_p: float = 0.0, + softmax_scale: float = None, + causal: bool = False, + return_attn_probs: bool = False, + out: torch.Tensor = None, +): + """ + Args: + q: (total_q, nheads, headdim) torch.float16, torch.bfloat16 + where total_q = total number of query tokens in the batch. + k: (total_k, nheads_k, headdim) torch.float16, torch.bfloat16 + where total_k = total number of key tokens in the batch. + v: (total_k, nheads_k, headdim) torch.float16, torch.bfloat16 + cu_seqlens_q: (batch_size + 1) torch.int32 + The cumulative sequence lengths of the sequences in the batch, used to index into q. + cu_seqlens_k: (batch_size + 1) torch.int32 + The cumulative sequence lengths of the sequences in the batch, used to index into kv. + max_seqlen_q: int + Maximum query sequence length in the batch. + max_seqlen_k: int + Maximum key sequence length in the batch. + dropout_p: float + Dropout probability. dropout_p should be set to 0.0 during evaluation + softmax_scale: float + The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim). + causal: bool + Whether to apply causal attention mask (e.g., for auto-regressive modeling). + return_attn_probs: bool + Whether to return the attention probabilities. This option is for testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling). + out: (total, nheads, headdim) torch.float16, torch.bfloat16 + Returns: + out: (total, nheads, headdim) torch.float16, torch.bfloat16 + """ + + + assert len(q.shape) == 3, "q.shape != [total_q, nheads, head_dim]" + assert len(k.shape) == 3, "k.shape != [total_k, nheads_k, head_dim]" + assert len(v.shape) == 3, "v.shape != [total_k, nheads_k, head_dim]" + assert len(cu_seqlens_q.shape) == 1, "cu_seqlens_q.shape != [batch_size+1]" + assert len(cu_seqlens_k.shape) == 1, "cu_seqlens_k.shape != [batch_size+1]" + + if return_attn_probs: + raise NotImplementedError("return_attn_probs not supported!") + atten_scale = softmax_scale + training = q.requires_grad + nheads = q.size(1) + nheads_k = k.size(1) + if training: + raise NotImplementedError("not support training!") + else: # 推理支持group query attention + assert nheads % nheads_k == 0 + return ixinfer_flash_attn_unpad( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + atten_scale, + out=out, + ) + + +def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: + """torch.repeat_interleave(x, dim=2, repeats=n_rep)""" + if len(x.shape) == 4: + batch, seq_len, n_kv_heads, head_dim = x.shape + elif len(x.shape) == 3: + tokens, n_kv_heads, head_dim = x.shape + if n_rep == 1: + return x + if len(x.shape) == 4: + return ( + x[:, :, :, None, :] + .expand(batch, seq_len, n_kv_heads, n_rep, head_dim) + .reshape(batch, seq_len, n_kv_heads * n_rep, head_dim) + ) + elif len(x.shape) == 3: + return ( + x[:, :, None, :] + .expand(tokens, n_kv_heads, n_rep, head_dim) + .reshape(tokens, n_kv_heads * n_rep, head_dim) + ) + + +def mha(q, k, v, atten_scale, is_causal): + q = q.permute(0, 2, 1, 3).contiguous() # batch num_head seq_len head_dim + k = k.permute(0, 2, 1, 3).contiguous() + v = v.permute(0, 2, 1, 3).contiguous() + + # 2. q*kt softmax + scores_qk = torch.matmul(q.float(), k.float().transpose(-2, -1)) * atten_scale + q_seq_len = q.size(2) + kv_seq_len = k.size(2) + if is_causal: + # Create attention mask. + attn_mask = torch.triu( + torch.ones(q_seq_len, kv_seq_len, dtype=torch.int), diagonal=1 + ) + attn_mask = attn_mask.to(dtype=torch.int, device="cuda") + else: + attn_mask = None + # softmax + # print(scores_qk.shape,attn_mask.shape) + if attn_mask is not None: + # print(scores_qk.shape,attn_mask.shape) + scores_qk = scores_qk + attn_mask * (-100000) + scores_qk = torch.nn.functional.softmax(scores_qk, dim=-1) + # 3. x = qk_scores * v + scores_v = torch.matmul(scores_qk, v.float()) + scores_v = scores_v.half() + return scores_v.permute(0, 2, 1, 3).contiguous() + + +def ref_flash_attn_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dropout_p: float = 0.0, + softmax_scale: float = None, + causal: bool = False, + return_attn_probs: bool = False, +): + if return_attn_probs: + raise NotImplementedError("return_attn_probs not supported!") + head_num = q.size(2) + head_num_kv = k.size(2) + if head_num != head_num_kv: + k = repeat_kv(k, head_num // head_num_kv) # [0,0,0,0,1,1,1,1,2,2,2,2] GROUP + v = repeat_kv(v, head_num // head_num_kv) + output_pt = mha(q, k, v, softmax_scale, causal) + return output_pt + + +def flash_attn_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dropout_p: float = 0.0, + softmax_scale: float = None, + causal: bool = False, + return_attn_probs: bool = False, +): + """ + Args: + q: (batch_size, seqlen, nheads, headdim) torch.float16, torch.bfloat16 + k: (batch_size, seqlen, nheads_k, headdim) torch.float16, torch.bfloat16 + v: (batch_size, seqlen, nheads_k, headdim) torch.float16, torch.bfloat16 + dropout_p: float + Dropout probability. dropout_p should be set to 0.0 during evaluation + softmax_scale: float + The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim). + causal: bool + Whether to apply causal attention mask (e.g., for auto-regressive modeling). + return_attn_probs: bool + Whether to return the attention probabilities. This option is for testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling). + Returns: + Tensor: (total, nheads, headdim) torch.float16, torch.bfloat16 + """ + + if return_attn_probs: + raise NotImplementedError("return_attn_probs not supported!") + atten_scale = softmax_scale + training = q.requires_grad + + q_dim = q.dim() + assert q_dim == 4 + + batch_size, max_seqlen_q, nheads, head_dim = q.shape + _, max_seqlen_k, nheads_k, head_dim_k = k.shape + assert head_dim == head_dim_k + if training: + raise NotImplementedError("not support training!") + else: # 推理支持group query attention + assert nheads % nheads_k == 0 + + q = q.view(batch_size * max_seqlen_q, nheads, head_dim) + k = k.view(batch_size * max_seqlen_k, nheads_k, head_dim) + v = v.view(batch_size * max_seqlen_k, nheads_k, head_dim) + + cu_seqlens_q = torch.ones([batch_size + 1]) * max_seqlen_q + cu_seqlens_q[0] = 0 + cu_seqlens_k = torch.ones([batch_size + 1]) * max_seqlen_k + cu_seqlens_k[0] = 0 + cu_seqlens_q = cu_seqlens_q.cuda().int() + cu_seqlens_k = cu_seqlens_k.cuda().int() + cu_seqlens_q = torch.cumsum(cu_seqlens_q, dim=0).int() + cu_seqlens_k = torch.cumsum(cu_seqlens_k, dim=0).int() + output = ixinfer_flash_attn_unpad( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + causal, + atten_scale, + sqrt_alibi=False, + alibi_slopes=None, + ) + return output.view(batch_size, max_seqlen_q, nheads, head_dim) diff --git a/ixformer_sdk/inference/functions/fused_rope.py b/ixformer_sdk/inference/functions/fused_rope.py new file mode 100644 index 00000000..5c95200b --- /dev/null +++ b/ixformer_sdk/inference/functions/fused_rope.py @@ -0,0 +1,76 @@ +from typing import List, Tuple, Union + +import ixformer._C as ops +import torch + +# adding by xuelu.peng 20240417 +# from https://github.com/NVIDIA/apex/blob/master/apex/transformer/functional/fused_rope.py#L59 +__all__ = ["fused_apply_rotary_pos_emb", "ref_fused_apply_rotary_pos_emb"] + +# Copied from Megatron-Core for testing. +# https://github.com/NVIDIA/Megatron-LM/blob/5f2877d85cb26e47ce6dcdae4b80adf376abf4e8/megatron/core/models/common/embeddings/rotary_pos_embedding.py#L139 +def apply_rotary_pos_emb(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Apply rotary positional embedding to input tensor T. + + check https://kexue.fm/archives/8265 for detailed formulas + + Arguments: + t (Tensor): Input tensor T is of shape [seq_length, ... , dim] + freqs (Tensor): Rotary Positional embedding tensor freq is of shape [seq_length, ..., dim] + + Returns: + Tensor: The input tensor after applying RoPE + """ + rot_dim = freqs.shape[-1] + + # ideally t_pass is empty so rotary pos embedding is applied to all tensor t + t, t_pass = t[..., :rot_dim], t[..., rot_dim:] + + # first part is cosine component + # second part is sine component, need to change signs with _rotate_half method + cos_ = torch.cos(freqs).to(t.dtype) + sin_ = torch.sin(freqs).to(t.dtype) + + t = (t * cos_) + (_rotate_half(t) * sin_) + return torch.cat((t, t_pass), dim=-1) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Change sign so the last dimension becomes [-odd, +even] + + Arguments: + x (Tensor): Input tensor + + Returns: + Tensor: Tensor rotated half + """ + + x1, x2 = torch.chunk(x, 2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def ref_fused_apply_rotary_pos_emb( + t: torch.Tensor, freqs: torch.Tensor, transpose_output_memory: bool = False +): + output_unfused = apply_rotary_pos_emb(t, freqs) + return output_unfused + + +def fused_apply_rotary_pos_emb( + t: torch.Tensor, + freqs: torch.Tensor, + transpose_output_memory: bool = False, +) -> torch.Tensor: + + """ + Args: + t: (sequence length,batch size,head num,head_dim) torch.float16, torch.bfloat16, torch.float32 + freqs: (sequence length,1 ,1, head_dim) torch.float32 + transpose_output_memory: bool + Default to False. Whether to transpose the 's' and 'b' dimension of the output's underlying memory format. This is very helpful when you want to get a contiguous tensor after calling `output.transpose(0, 1)`. + + Returns: + Tensor: (sequence length,batch size,head num,head_dim) torch.float16, torch.bfloat16, torch.float32 + """ + output = ops.train.fused_rope_forward(t, freqs, transpose_output_memory) + return output diff --git a/ixformer_sdk/inference/functions/gemv.py b/ixformer_sdk/inference/functions/gemv.py new file mode 100644 index 00000000..ba3344d1 --- /dev/null +++ b/ixformer_sdk/inference/functions/gemv.py @@ -0,0 +1,48 @@ +import os +from typing import Union + +import ixformer._C as ops +import torch + +__all__ = ["gemv", "ref_gemv"] + + +def ref_gemv(x: torch.Tensor, A: torch.Tensor, gemv_max_batch: int = 1): + output = torch.nn.functional.linear(x, A) + return output + + +def gemv_conditions(input, weight, gemv_max_batch): + # gemv 使用的条件 input:[m,k] weight:[n,k] + # 1. m<=gemv_max_batch + # 2. k%2==0 n%2==0 + # 3. bias is None + input = input.view(-1, input.shape[-1]) + weight = weight.view(-1, weight.shape[-1]) + m = input.shape[0] + k = input.shape[1] + n = weight.shape[0] + if m <= gemv_max_batch and k % 2 == 0 and n % 2 == 0: + return True + return False + + +def gemv(x: torch.Tensor, A: torch.Tensor, gemv_max_batch: int = 1): + + """ + Args: + x: (..., k) torch.float16, torch.bfloat16 + A: (n,k) torch.float16, torch.bfloat16, torch.float32 + gemv_max_batch: int + 用于是否满足gemv使用条件的判断,目前只支持到1 + Returns: + Tensor: (..., n) torch.float16, torch.bfloat16 + """ + disable_infer_gemm_ex = os.getenv("DISABLE_INFER_GEMM_EX", "0") + use_gemv = gemv_conditions(x, A, gemv_max_batch) and disable_infer_gemm_ex != "1" + assert use_gemv == True + output_shape = list(x.shape) + output_shape[-1] = A.shape[0] + output = x.new_empty(output_shape) + output = ops.infer.linear_ex(x, A, None, output) + return output diff --git a/ixformer_sdk/inference/functions/groupnorm.py b/ixformer_sdk/inference/functions/groupnorm.py new file mode 100644 index 00000000..fa7e2efc --- /dev/null +++ b/ixformer_sdk/inference/functions/groupnorm.py @@ -0,0 +1,120 @@ +from typing import List, Tuple, Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +import ixformer + +__all__ = [ + "group_norm", + "ref_group_norm", + "ref_fused_group_norm_silu", + "fused_group_norm_silu", + "ref_fused_group_norm_silu_nhwc", + "fused_group_norm_silu_nhwc" +] +def is_channels_last(ten): + return torch._prims_common.suggest_memory_format(ten) == torch.channels_last + +def ref_group_norm(input, num_groups, weight, bias, eps): + output = torch.nn.functional.group_norm(input, num_groups, weight, bias, eps) + return output + +#group_norm官方接口，如果input是nhwc(channel_last),输出则不是channel_last,而是nchw；如果input是nchw，那么输出也是nchw +def group_norm( + input: torch.Tensor, + num_groups: int, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float = 1e-05, + +): + """ + Args: + input: (n,c,h,w) or (n,c,h) or (n,h,w,c) torch.float16 + "contiguous_format":(n,c,h,w) or (n,c,h) "channels_last": (n,h,w,c) + num_groups: int + weight: (c) torch.float16 + bias: (c) torch.float16 + eps: float + Returns: + Tensor: (n,c,h,w) torch.float16 + """ + + is_nhwc=is_channels_last(input) + out = ops.infer.groupnorm(input, num_groups, weight, bias, eps, is_nhwc, 0) + if is_nhwc: + out=out.permute(0,3,1,2).contiguous() + return out +def ref_fused_group_norm_silu_nhwc( + input: torch.Tensor, + num_groups: int, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float = 1e-05, + act_type:int = 0 +): + output = torch.nn.functional.group_norm(input.permute(0,3,1,2).contiguous(), num_groups, weight, bias, eps) + if act_type: + output = output * torch.sigmoid(output) + output = output.permute(0,2,3,1).contiguous() + return output +#为了减少permute/contiguous,新接口支持输入输出都是nhwc的，融合silu +def fused_group_norm_silu_nhwc( + input: torch.Tensor, + num_groups: int, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float = 1e-05, + act_type:int = 0 +): + """ + Args: + input: (n,h,w,c) torch.float16 + num_groups: int + weight: (c) torch.float16 + bias: (c) torch.float16 + eps: float + act_type: int + 0 or 1,if act_type=1, silu + Returns: + Tensor: (n,h,w,c) torch.float16 + + """ + out = ops.infer.groupnorm(input, num_groups, weight, bias, eps, True, act_type) + return out + +def ref_fused_group_norm_silu( + input: torch.Tensor, + num_groups: int, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float = 1e-05, +): + output = torch.nn.functional.group_norm(input, num_groups, weight, bias, eps) + output = output * torch.sigmoid(output) + return output + + +def fused_group_norm_silu( + input: torch.Tensor, + num_groups: int, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float = 1e-05, +): + + """ + Args: + input: (n,c,h,w) or (n,c,h) torch.float16 + num_groups: int + weight: (c) torch.float16 + bias: (c) torch.float16 + eps: float + Returns: + output: (n,c,h,w) or (n,c,h) torch.float16 + """ + + out = ops.infer.groupnorm(input, num_groups, weight, bias, eps, False, 1) + return out diff --git a/ixformer_sdk/inference/functions/i8w8o32.py b/ixformer_sdk/inference/functions/i8w8o32.py new file mode 100644 index 00000000..ba022c4a --- /dev/null +++ b/ixformer_sdk/inference/functions/i8w8o32.py @@ -0,0 +1,32 @@ +import os +from typing import Union + +import ixformer._C as ops +import torch + +__all__ = ["i8w8o32", "ref_i8w8o32"] + + +def ref_i8w8o32(input: torch.Tensor, weight: torch.Tensor): + output = torch.nn.functional.linear(input.float(), weight.float()).int() + return output + + +def i8w8o32(input: torch.Tensor, weight: torch.Tensor): + + """ + Args: + input: (bs, ic) torch.int8 + weight: (oc, ic) torch.int8 + Returns: + Tensor: (bs, oc)) torch.int32 + """ + if not torch.is_tensor(input): + raise RuntimeError("Not impl.") + output_shape = list(input.shape) + output_shape[-1] = weight.size(0) + output = torch.empty(output_shape, dtype=torch.int32, device=input.device) + ic_dim = input.size(-1) + input = input.view(-1, ic_dim) + ops.infer.linear_i8w8o32(input.view(-1, ic_dim), weight, output) + return output diff --git a/ixformer_sdk/inference/functions/layernorm.py b/ixformer_sdk/inference/functions/layernorm.py new file mode 100644 index 00000000..6aadc095 --- /dev/null +++ b/ixformer_sdk/inference/functions/layernorm.py @@ -0,0 +1,429 @@ +from typing import List, Tuple, Union + +import ixformer._C as ops +import torch + +__all__ = [ + "layer_norm", + "ref_layer_norm", + "residual_layer_norm", + "ref_residual_layer_norm", + "ref_residual_layer_norm_bias_alpha", + "residual_layer_norm_bias_alpha", + "ref_layer_norm_2sb_fused", + "layer_norm_2sb_fused", +] + + +def ref_layer_norm( + input: torch.Tensor, + normalized_shape: List[int], + weight: torch.Tensor, + bias: torch.Tensor, + eps: float = 1e-5, + output: torch.Tensor = None, +): + + if weight is None or bias is None or weight.dim() > 1 or bias.dim() > 1: + raise NotImplementedError( + "layer_norm only support weight.dim() ==1 and bias.dim()==1!" + ) + if normalized_shape == None: + norm_size = weight.size(-1) + normalized_shape = [norm_size] + else: + if ( + isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple) + ) and len(normalized_shape) == 1: + norm_size = normalized_shape[0] + else: + raise ValueError( + f"layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}" + ) + if norm_size != weight.size(-1): + raise ValueError(f"layer_norm(): argument 'norm_size' must == weight.size(-1)") + + norm_out = torch.nn.functional.layer_norm( + input, normalized_shape, weight, bias, eps=eps + ) + if output is not None: + assert output.shape == norm_out.shape + output.copy_(norm_out) + else: + output = norm_out + + return output + + +def layer_norm( + input: torch.Tensor, + normalized_shape: List[int], + weight: torch.Tensor, + bias: torch.Tensor, + eps: float = 1e-5, + output: torch.Tensor = None, +): + """ + This function is deprecated, please use residual_layer_norm. + 等价实现： + torch.nn.functional.layer_norm( input, normalized_shape, weight, bias, eps=0.000001) + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + normalized_shape: list[int] + weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + eps: float32 + Returns: + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + """ + if weight is None or bias is None or weight.dim() > 1 or bias.dim() > 1: + raise NotImplementedError( + "layer_norm only support weight.dim() ==1 and bias.dim()==1!" + ) + if normalized_shape == None: + norm_size = weight.size(-1) + normalized_shape = [norm_size] + else: + if ( + isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple) + ) and len(normalized_shape) == 1: + norm_size = normalized_shape[0] + else: + raise ValueError( + f"layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}" + ) + if norm_size != weight.size(-1): + raise ValueError(f"layer_norm(): argument 'norm_size' must == weight.size(-1)") + if output is None: + output = torch.empty_like(input) + ops.infer.layer_norm(input, weight, bias, None, output, eps) + return output + + +def ref_residual_layer_norm( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + residual_bias: torch.Tensor = None, + eps: float = 1e-5, + output: torch.Tensor = None, + residual_output: torch.Tensor = None, +): + normalized_shape = [weight.size(-1)] + + if residual_bias is not None: + input = input + residual_bias + + if residual is not None: + residual_output = torch.add(input, residual, out=residual_output) + input = residual_output + + norm_out = torch.nn.functional.layer_norm( + input, normalized_shape, weight, bias, eps=eps + ) + + if output is None: + output = norm_out + else: + output.copy_(norm_out) + + return output, residual_output + + +def residual_layer_norm( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + residual_bias: torch.Tensor = None, + eps: float = 1e-5, + output: torch.Tensor = None, + residual_output: torch.Tensor = None, +): + """ + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + eps: float32 + Returns: + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on input. + residual_output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on residual. + """ + if residual is None: + if output is None: + output = torch.empty(input.shape, device=input.device, dtype=input.dtype) + ops.infer.layer_norm(input, weight, bias, residual_bias, output, eps) + else: + ops.infer.residual_layer_norm( + input, + residual, + weight, + bias, + residual_bias, + output, + residual_output, + 1.0, + eps, + False, + ) + residual_output = residual_output if residual_output is not None else residual + output = output if output is not None else input + + return output, residual_output + + +def ref_residual_layer_norm_bias_alpha( + input: torch.Tensor, + normalized_shape: List[int], + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor, + residual_bias: torch.Tensor = None, + alpha: float = 1.0, + eps: float = 1e-5, + is_post_ln=False, +): + if ( + weight is None + or bias is None + or residual is None + or weight.dim() > 1 + or bias.dim() > 1 + ): + raise NotImplementedError( + "residual_layer_norm only support weight.dim() ==1 and bias.dim()==1!" + ) + if normalized_shape == None: + norm_size = weight.size(-1) + normalized_shape = [norm_size] + else: + if ( + isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple) + ) and len(normalized_shape) == 1: + norm_size = normalized_shape[0] + else: + raise ValueError( + f"residual_layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}" + ) + + if norm_size != weight.size(-1): + raise ValueError( + f"residual_layer_norm(): argument 'norm_size' must == weight.size(-1)" + ) + dtype = input.dtype + if residual_bias is None: + x = input.float() + residual.float() * alpha + else: + x = input.float() + residual.float() * alpha + residual_bias.float() + + y = torch.nn.functional.layer_norm( + x.to(dtype), normalized_shape, weight, bias, eps=eps + ) + + if is_post_ln: + return y, y + else: + return y, x.to(dtype) + + +def residual_layer_norm_bias_alpha( + input: torch.Tensor, + normalized_shape: List[int], + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor, + residual_bias: torch.Tensor = None, + alpha: float = 1.0, + eps: float = 1e-5, + is_post_ln=False, +): + """ + 等价实现： + residual = input + residual.float() * alpha + residual_bias + output = torch.nn.functional.layer_norm( + residual, normalized_shape, weight, bias, eps=eps + ) + residual = output if is_post_ln else residual + + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + normalized_shape list[int] + weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + alpha: float32 + eps: float32 + is_post_ln: bool + Returns: + input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + Inplace operation will be performed on residual and input. + """ + if ( + weight is None + or bias is None + or residual is None + or weight.dim() > 1 + or bias.dim() > 1 + ): + raise NotImplementedError( + "residual_layer_norm only support weight.dim() ==1 and bias.dim()==1!" + ) + if normalized_shape == None: + norm_size = weight.size(-1) + normalized_shape = [norm_size] + else: + if ( + isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple) + ) and len(normalized_shape) == 1: + norm_size = normalized_shape[0] + else: + raise ValueError( + f"residual_layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}" + ) + + if norm_size != weight.size(-1): + raise ValueError( + f"residual_layer_norm(): argument 'norm_size' must == weight.size(-1)" + ) + + ops.infer.residual_layer_norm( + input, residual, weight, bias, residual_bias, None, None, alpha, eps, is_post_ln + ) + + return input, residual + + +def ref_layer_norm_2sb_fused( + input: torch.Tensor, + normalized_shape: List[int], + weight1: torch.Tensor, + bias1: torch.Tensor, + weight2: torch.Tensor, + bias2: torch.Tensor, + eps: float = 1e-5, +): + assert input.shape[-1] <= 16384 + if not (input.dtype == torch.float16 or input.dtype == torch.bfloat16): + raise NotImplementedError( + "layer_norm_2sb() only support data format of float16 or bfloat16 now!" + ) + if ( + weight1 is None + or bias1 is None + or weight2 is None + or bias2 is None + or weight1.dim() > 1 + or bias1.dim() > 1 + or weight2.dim() > 1 + or bias2.dim() > 1 + ): + raise NotImplementedError( + "layer_norm_2sb only support weight1.dim() ==1, bias1.dim()==1, weight2.dim() ==1 and bias2.dim()==1 !" + ) + if normalized_shape == None: + norm_size = weight1.size(-1) + normalized_shape = [norm_size] + else: + if ( + isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple) + ) and len(normalized_shape) == 1: + norm_size = normalized_shape[0] + else: + raise ValueError( + f"layer_norm_2sb(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}" + ) + + if norm_size != weight1.size(-1) or norm_size != weight2.size(-1): + raise ValueError( + f"layer_norm_2sb(): argument 'norm_size' must == weight.size(-1)" + ) + + output1 = torch.nn.functional.layer_norm( + input, normalized_shape, weight1, bias1, eps=eps + ) + + output2 = torch.nn.functional.layer_norm( + input, normalized_shape, weight2, bias2, eps=eps + ) + + return output1, output2 + + +def layer_norm_2sb_fused( + input: torch.Tensor, + normalized_shape: List[int], + weight1: torch.Tensor, + bias1: torch.Tensor, + weight2: torch.Tensor, + bias2: torch.Tensor, + eps: float = 1e-5, +): + """ + 等价实现： + output1 = torch.nn.functional.layer_norm( + input, normalized_shape, weight1, bias1, eps=eps + ) + output2 = torch.nn.functional.layer_norm( + input, normalized_shape, weight2, bias2, eps=eps + ) + + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16 + normalized_shape list[int] + weight1: (hidden_size) torch.float16, torch.bfloat16 + bias1: (hidden_size) torch.float16, torch.bfloat16 + weight2: (hidden_size) torch.float16, torch.bfloat16 + bias2: (hidden_size) torch.float16, torch.bfloat16 + eps: float32 + Returns: + output1: (..., hidden_size) torch.float16, torch.bfloat16 + output2: (..., hidden_size) torch.float16, torch.bfloat16 + """ + assert input.shape[-1] <= 16384 + if not (input.dtype == torch.float16 or input.dtype == torch.bfloat16): + raise NotImplementedError( + "layer_norm_2sb() only support data format of float16 or bfloat16 now!" + ) + if ( + weight1 is None + or bias1 is None + or weight2 is None + or bias2 is None + or weight1.dim() > 1 + or bias1.dim() > 1 + or weight2.dim() > 1 + or bias2.dim() > 1 + ): + raise NotImplementedError( + "layer_norm_2sb only support weight1.dim() ==1, bias1.dim()==1, weight2.dim() ==1 and bias2.dim()==1 !" + ) + if normalized_shape == None: + norm_size = weight1.size(-1) + normalized_shape = [norm_size] + else: + if ( + isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple) + ) and len(normalized_shape) == 1: + norm_size = normalized_shape[0] + else: + raise ValueError( + f"layer_norm_2sb(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}" + ) + + if norm_size != weight1.size(-1) or norm_size != weight2.size(-1): + raise ValueError( + f"layer_norm_2sb(): argument 'norm_size' must == weight.size(-1)" + ) + output1 = torch.empty_like(input) + output2 = torch.empty_like(input) + ops.infer.layer_norm_2sb( + input, weight1, bias1, weight2, bias2, eps, output1, output2 + ) + + return output1, output2 diff --git a/ixformer_sdk/inference/functions/lightllm.py b/ixformer_sdk/inference/functions/lightllm.py new file mode 100644 index 00000000..3033e232 --- /dev/null +++ b/ixformer_sdk/inference/functions/lightllm.py @@ -0,0 +1,277 @@ +from typing import Union + +import ixformer._C as ops +import torch + +__all__ = [ + "lightllm_tokenattention", + "ref_lightllm_tokenattention", + "lightllm_destindex_copy_kv", + "ref_lightllm_destindex_copy_kv", + "lightllm_apply_penalty", + "ref_lightllm_apply_penalty", + "lightllm_glm2_rope", + "ref_lightllm_glm2_rope", +] + + +def ref_lightllm_glm2_rope( + x: torch.Tensor, # tokens,head_num,head_dim + cos: torch.Tensor, # tokens,rotdim + sin: torch.Tensor, +): + num_tokens, _, rot_dim = list(cos.shape) + head_num = x.shape[1] + x12 = x[:, :, : rot_dim * 2] + x3 = x[:, :, rot_dim * 2 :] + x12 = x12.reshape(num_tokens, head_num, rot_dim, 2) + x1 = x12[:, :, :, 0] + x2 = x12[:, :, :, 1] + + # out0 = q0 * cos - q1 * sin + # out1 = q0 * sin + q1 * cos + # q1, q2 是沿着 head_dim维度，交叉取值的 + + q1 = x1 * cos - x2 * sin + q2 = x2 * cos + x1 * sin + + q12 = torch.stack([q1, q2], dim=-1) + q12 = q12.reshape(num_tokens, head_num, -1) + x_pytorch = torch.cat([q12, x3], dim=-1) + return x_pytorch + + +def lightllm_glm2_rope( + x: torch.Tensor, # tokens,head_num,head_dim + cos: torch.Tensor, # tokens,rotdim + sin: torch.Tensor, +): + """ + Args: + x: (num_tokens, head_num, head_dim) torch.half + cos: (num_tokens,1,head_dim//2//2) torch.half + sin: (num_tokens,1,head_dim//2//2) torch.half + Returns: + x: (num_tokens, head_num, head_dim) torch.half + + """ + if isinstance(x, torch.Tensor): + ops.infer.lightllm_glm2_rope(x, cos, sin) + return x + else: + raise NotImplementedError() + + +def ref_lightllm_apply_penalty( + Logits: torch.Tensor, + presence_penalty: torch.Tensor, + freqency_penalty: torch.Tensor, + p_token_ids: torch.Tensor, + p_token_counts: torch.Tensor, + p_cumsum_seq_len: torch.Tensor, + p_max_len_in_batch: int, +): + batch_size = Logits.size(0) + output = Logits.clone() + for cur_batch in range(batch_size): + cur_freqency = freqency_penalty[cur_batch] + cur_presence = presence_penalty[cur_batch] + cur_batch_start_index = p_cumsum_seq_len[cur_batch] + cur_batch_end_index = p_cumsum_seq_len[cur_batch + 1] + for token_idx in range(cur_batch_start_index, cur_batch_end_index): + batch_ids = p_token_ids[token_idx] + batch_ids_count = p_token_counts[token_idx] + cur_logits = output[cur_batch][batch_ids] + + freq_logits = cur_logits - batch_ids_count * cur_freqency + pre_logits = freq_logits - cur_presence + # if token_idx==0: + # print(f"batch_ids {batch_ids} cur_logits {cur_logits} pre_logits {pre_logits}") + output[cur_batch][batch_ids] = pre_logits + return output + + +def lightllm_apply_penalty( + Logits: torch.Tensor, + presence_penalty: torch.Tensor, + freqency_penalty: torch.Tensor, + p_token_ids: torch.Tensor, + p_token_counts: torch.Tensor, + p_cumsum_seq_len: torch.Tensor, + p_max_len_in_batch: int, +): + """ + Args: + logits: (batch_size, vocab_size) torch.float + presence_penalty: (batch_size) torch.float + freqency_penalty: (batch_size) torch.float + p_token_ids: (num_tokens) torch.int + p_token_counts: (num_tokens) torch.int + p_cumsum_seq_len: (batch_size+1) torch.int + p_max_len_in_batch: int + 在一个batch中seq的最大长度 + Returns: + logits: (batch_size, vocab_size) torch.float + """ + if isinstance(Logits, torch.Tensor): + ops.infer.lightllm_apply_penalty( + Logits, + presence_penalty, + freqency_penalty, + p_token_ids, + p_token_counts, + p_cumsum_seq_len, + p_max_len_in_batch, + ) + return Logits + else: + raise NotImplementedError() + + +def ref_lightllm_destindex_copy_kv( + key_cache: torch.Tensor, + mem_idx: torch.Tensor, + output: torch.Tensor, +): + if key_cache.dim() != 3 or key_cache.size(-1) != 128: + raise NotImplementedError( + "lightllm_destindex_copy_kv only support key_cache.dim()==3 and head_size ==128 !" + ) + output[mem_idx.long()] = key_cache + return output + + +def lightllm_destindex_copy_kv( + key_cache: torch.Tensor, + mem_idx: torch.Tensor, + output: torch.Tensor, +): + """ + Args: + key_cache: (tokens, num_kv_heads, head_size) torch.half + 目前head_size 只支持128的情况 + mem_idx: (tokens) torch.int + output: (max_tokens, num_kv_heads, head_size) torch.half + Returns: + output: (max_tokens, num_kv_heads, head_size) torch.half + """ + + if key_cache.dim() != 3 or key_cache.size(-1) != 128: + raise NotImplementedError( + "lightllm_destindex_copy_kv only support key_cache.dim()==3 and head_size ==128 !" + ) + if isinstance(key_cache, torch.Tensor): + ops.infer.lightllm_destindex_copy_kv(key_cache, mem_idx, output) + else: + raise NotImplementedError() + return output + + +def ref_lightllm_tokenattention( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + reg_tokens: torch.Tensor, + b_req_idx: torch.Tensor, + b_seq_len: torch.Tensor, + scale: float, + max_context_len: int, +): + batch_size, tp_q_head_num_, head_dim_ = query.shape + tp_k_head_num_ = key_cache.size(-2) + sm_scale = scale + curbatch_max_context_len = max_context_len + tmp_k = torch.zeros( + (batch_size, tp_q_head_num_, curbatch_max_context_len, head_dim_), + dtype=query.dtype, + device="cuda", + ) + tmp_v = torch.zeros( + (batch_size, tp_q_head_num_, curbatch_max_context_len, head_dim_), + dtype=query.dtype, + device="cuda", + ) + mask = torch.ones([batch_size, 1, 1, curbatch_max_context_len]) + + kv_group_num = tp_q_head_num_ // tp_k_head_num_ + for cur_batch in range(batch_size): + cur_batch_req_idx = b_req_idx[cur_batch] + seq_len = b_seq_len[cur_batch] + mask[cur_batch, :, :, :seq_len] = 0 + # print(f"cur_batch {cur_batch}") + + for seq_idx in range(seq_len): + k_loc = reg_tokens[cur_batch_req_idx][seq_idx] + # print(k_loc) + for cur_head in range(tp_q_head_num_): + cur_kv_head = cur_head // kv_group_num + tmp_k[cur_batch, cur_head, seq_idx, :] = key_cache[k_loc][cur_kv_head] + tmp_v[cur_batch, cur_head, seq_idx, :] = value_cache[k_loc][cur_kv_head] + mask = mask.cuda() + # batch_size, self.tp_q_head_num_, 1, max_len_in_batch + attn_score = ( + torch.matmul( + query.view(batch_size, tp_q_head_num_, 1, head_dim_), + tmp_k.transpose(-1, -2), + ) + * sm_scale + ) + attn_score = attn_score + mask * -1000 + attn_score = torch.softmax(attn_score, dim=-1) + # batch_size, self.tp_q_head_num_, 1, head_dim + py_out = torch.matmul(attn_score.to(query.dtype), tmp_v).view( + batch_size, tp_q_head_num_, -1 + ) + return py_out + + +def lightllm_tokenattention( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + reg_tokens: torch.Tensor, + b_req_idx: torch.Tensor, + b_seq_len: torch.Tensor, + scale: float, + max_context_len: int, + partition: int, + output: torch.Tensor, +): + """ + Args: + query: (batch_size,head_num,head_dim) torch.float16, torch.bfloat16 + key_cache: (max_num_tokens, head_num_kv, head_dim) torch.float16, torch.bfloat16 + value_cache: (max_num_tokens, head_num_kv, head_dim) torch.float16, torch.bfloat16 + reg_tokens: (max_request,max_tokens) torch.int32 + 目前max_tokens只支持3080 + b_req_idx: (batch_size) torch.int32 + b_req_len: (batch_size) torch.int32 + scale: float + The scaling of QK^T before applying softmax. + max_context_len: int + b_seq_len.max() + partition: int + Returns: + output: (batch_size,head_num,head_dim) torch.float16, torch.bfloat16 + """ + _,max_tokens=reg_tokens.shape + if not max_tokens == 3080: + raise NotImplementedError( + "lightllm_tokenattention only support reg_tokens.size(-1)==3080" + ) + if isinstance(query, torch.Tensor): + ops.infer.lightllm_tokenattention( + query, + key_cache, + value_cache, + reg_tokens, + b_req_idx, + b_seq_len, + scale, + max_context_len, + partition, + output, + ) + else: + raise NotImplementedError() + return output diff --git a/ixformer_sdk/inference/functions/linalg.py b/ixformer_sdk/inference/functions/linalg.py new file mode 100644 index 00000000..5566684f --- /dev/null +++ b/ixformer_sdk/inference/functions/linalg.py @@ -0,0 +1,50 @@ +from typing import Union + +import ixformer._C as ops +import torch + +__all__ = ["solve", "ref_slove"] + + +def ref_slove( + A: torch.Tensor, B: torch.Tensor, *, left: bool = True, out: torch.Tensor = None +): + out = torch.linalg.solve(A, B, left=left) + return out + + +def solve( + A: torch.Tensor, B: torch.Tensor, *, left: bool = True, out: torch.Tensor = None +): + """ + Args: + A: (..., n, n) torch.float + B: (..., n) or (..., n, k) or (n,...) or (n, k) or (n) torch.float + left: bool + whether to solve the system AX=B or XA=B. Default: True, 目前只支持left =True + out: (..., n, k) torch.float + Returns: + out: (..., n, k) torch.float + """ + + n = A.shape[-1] + batch_count = A.numel() // (n * n) + if B.dim() == 1: + k = 1 + elif B.dim() == 2: + if A.dim() > 2 and B.shape == (batch_count, n): + k = 1 + else: + nid = 0 if left else 1 + k = B.shape[nid ^ 1] + else: + k = B.size(B.dim() - 1 if left else B.dim() - 2) + + if n <= 64 and k <= 64 and left: + return ops.infer.solve(A, B, left) + else: + device = A.device + cpu_A = A.cpu() + cpu_B = B.cpu() + cpu_res = ref_slove(A=cpu_A, B=cpu_B, left=left) + return cpu_res.to(device) diff --git a/ixformer_sdk/inference/functions/linear.py b/ixformer_sdk/inference/functions/linear.py new file mode 100644 index 00000000..507eef34 --- /dev/null +++ b/ixformer_sdk/inference/functions/linear.py @@ -0,0 +1,122 @@ +import os +from typing import Union + +import ixformer._C as ops +import torch + +__all__ = ["linear", "ref_linear", "mixed_type_linear", "ref_mixed_type_linear"] + + +def ref_linear( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + act_type=-1, +): + output = torch.nn.functional.linear(input, weight, bias) + if act_type == -1: + act_fn = torch.nn.Identity() + elif act_type == 3: + act_fn = torch.nn.GELU() + elif act_type == 4: + act_fn = torch.nn.ReLU() + elif act_type == 12: + act_fn = torch.nn.SiLU() + else: + raise KeyError("act_type not supported") + output = act_fn(output) + return output + + +def gemv_conditions(input, weight, bias, gemv_max_batch): + # gemv 使用的条件 input:[m,k] weight:[n,k] + # 1. m<=gemv_max_batch + # 2. k%32==0 n%2==0 + # 3. bias is None + input = input.view(-1, input.shape[-1]) + weight = weight.view(-1, weight.shape[-1]) + m = input.shape[0] + k = input.shape[1] + n = weight.shape[0] + if bias is None and m <= gemv_max_batch and k % 32 == 0 and n % 2 == 0: + return True + return False + + +def linear( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + output: torch.Tensor = None, + persistent: bool = False, + act_type : int = -1, +): + """ + Args: + input: (...,k) torch.float16, torch.bfloat16 + weight: (n, k) torch.float16, torch.bfloat16 + bias: (n) torch.float16, torch.bfloat16 + output: (...,n) torch.float16, torch.bfloat16 + persistent: bool + 是否限制 Gemm Kernel 的 Block 数量 + Returns: + output: (...,n) torch.float16, torch.bfloat16 + """ + if not input.is_contiguous(): + input = input.contiguous() + if not weight.is_contiguous(): + weight = weight.contiguous() + use_gemv = True + gemv_max_batch = 1 + disable_infer_gemm_ex = os.getenv("DISABLE_INFER_GEMM_EX", "0") + use_gemv = ( + use_gemv + and gemv_conditions(input, weight, bias, gemv_max_batch) + and disable_infer_gemm_ex != "1" + ) + + if output is None: + output_shape = list(input.shape) + output_shape[-1] = weight.shape[0] + output = input.new_empty(output_shape) + + if not use_gemv: + output = ops.infer.linear(input, weight, act_type, bias, output, persistent) + else: + output = ops.infer.linear_ex(input, weight, bias, output) + return output + + +def ref_mixed_type_linear( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + output: torch.Tensor = None, + persistent=False, # TODO: support persistent +): + input = input.to(weight.dtype) + if bias: + bias = bias.to(weight.dtype) + output = torch.nn.functional.linear(input, weight, bias) + return output + + +def mixed_type_linear( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + output: torch.Tensor = None, + persistent=False, # TODO: support persistent +): + """ + Args: + input: (...,k) torch.half, torch.bfloat16 + weight: (m, k) torch.float32 + bias: not supported + output: (...,m) torch.float32 + persistent: bool + Returns: + output: (...,m) torch.float32 + """ + output = ops.infer.mixed_type_linear(input, weight, bias, output) + return output diff --git a/ixformer_sdk/inference/functions/lmdeploy.py b/ixformer_sdk/inference/functions/lmdeploy.py new file mode 100644 index 00000000..7a25f719 --- /dev/null +++ b/ixformer_sdk/inference/functions/lmdeploy.py @@ -0,0 +1,212 @@ +import math +from typing import Literal, Optional, Union + +import ixformer._C as ops +import ixformer._C._functions as CF +import torch + +from ixformer.core import config + +from .linear import linear +from .paged_attention import paged_attention as paged_attention_ixformer_impl + +__all__ = [ + "ref_lmdeploy_paged_attention", + "lmdeploy_paged_attention", +] + +weak_ref_tensor = ops.infer.weak_ref_tensor + + +def ref_lmdeploy_paged_attention( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + softcap: float = 0.0, + window_left: int = -1, + window_right: int = -1, + use_sqrt_alibi: bool = False, + quant_type: int = 0, + is_bbhh: bool = False, +): + assert window_right in [-1, 0] + + if is_bbhh: + key_cache = key_cache.permute(0, 2, 1, 3).contiguous() + value_cache = value_cache.permute(0, 2, 1, 3).contiguous() + + def get_alibi_mask(num_heads, seqlen, device, dtype): + x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1) + y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1) + offsets = -(y - x).view(1, 1, seqlen) + return offsets + + def ref_masked_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + query = query * scale + dtype = query.dtype + device = query.device + query = query.to(torch.float32) + key = key.to(torch.float32) + value = value.to(torch.float32) + attn = torch.einsum("qhd,khd->hqk", query, key) + if attn_mask is not None: + attn_mask = attn_mask + attn = attn + attn_mask + attn = torch.softmax(attn, dim=-1) + out = torch.einsum("hqk,khd->qhd", attn, value) + out = out.to(device).to(dtype) + return out + + head_size = query.shape[-1] + num_query_heads = query.shape[1] + num_kv_heads = value_cache.shape[1] + num_input_tokens = query.shape[0] + + num_q_per_kv = num_query_heads // num_kv_heads + slopes = ( + alibi_slopes.view(num_query_heads, 1, 1) + if alibi_slopes is not None + else alibi_slopes + ) + + for i in range(num_input_tokens): + q = query[i].unsqueeze(0) + block_table = block_tables[i] + context_len = int(context_lens[i]) + + keys = [] + values = [] + for j in range(context_len): + block_number = int(block_table[j // block_size]) + block_offset = j % block_size + + k = key_cache[block_number, :, block_offset, :] + keys.append(k) + + v = value_cache[block_number, :, block_offset, :] + values.append(v) + keys = torch.stack(keys, dim=0) + values = torch.stack(values, dim=0) + if num_q_per_kv > 1: + keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1) + values = torch.repeat_interleave(values, num_q_per_kv, dim=1) + if alibi_slopes is not None: + offsets = get_alibi_mask( + num_query_heads, context_len, output.device, output.dtype + ) + mask = offsets * slopes + mask = mask.to(output.dtype) + if window_left != -1: + index = torch.ones_like(mask, dtype=torch.int32, device=mask.device) + index[:, :, (context_len - 1 - window_left) :] = 0 + index = index.bool() + mask.masked_fill_(index, float("-inf")) + else: + if window_left != -1: + mask = torch.zeros([1, 1, context_len], dtype=q.dtype, device=q.device) + index = torch.ones_like(mask, dtype=torch.int32, device=mask.device) + index[:, :, (context_len - 1 - window_left) :] = 0 + index = index.bool() + mask.masked_fill_(index, float("-inf")) + else: + mask = None + + out = ref_masked_attention( + q, + keys, + values, + scale, + mask, + ) + out = out.view(num_query_heads, head_size) + if softcap != 0.0: + out = softcap * torch.tanh(out / softcap) + output[i].copy_(out, non_blocking=True) + + return output + + +def lmdeploy_paged_attention( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + softcap: float = 0.0, + causal: bool = True, + window_left: int = -1, + window_right: int = -1, + use_cuda_graph: bool = False, + use_sqrt_alibi: bool = False, + quant_type: int = 0, + is_bbhh: bool = False, +): + """ + is_bbhh = False key_cache, value_cache: [num_blocks, block_size, num_kv_heads, head_size] + is_bbhh = True key_cache, value_cache: [num_blocks, num_kv_heads, block_size, head_size] + + is_bbhh = False + Arguments: + query: [torch.half, torch.bfloat16] [num_tokens, num_heads, head_size] + key_cache: [torch.half, torch.bfloat16] [num_blocks, num_kv_heads, block_size, head_size] + value_cache: [torch.half, torch.bfloat16] [num_blocks, num_kv_heads, block_size, head_size] + num_kv_heads: int + scale: float + block_tables: [torch.int64] [num_tokens, max_num_blocks_per_seq] + context_lens: [torch.int32] [num_tokens] + block_size: int + max_context_len: int + alibi_slopes: [torch.float32] [num_heads] + softcap: float + causal: bool + window_left: int + window_right: int + use_sqrt_alibi: bool: False + Return: + output: [torch.half, torch.bfloat16] [num_tokens, num_heads, head_size] + """ + + ops.infer.lmdeploy_paged_attention( + output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + context_lens, + block_size, + max_context_len, + alibi_slopes, + causal, + window_left, + window_right, + softcap, + use_cuda_graph, + use_sqrt_alibi, + is_bbhh, + quant_type, + ) + return output + + # lmdeploy_paged_attention = lmdeploy_paged_attention_ixinfer diff --git a/ixformer_sdk/inference/functions/marlin.py b/ixformer_sdk/inference/functions/marlin.py new file mode 100644 index 00000000..07e66d19 --- /dev/null +++ b/ixformer_sdk/inference/functions/marlin.py @@ -0,0 +1,234 @@ +import ixformer._C as ops +import torch + +__all__ = [ + "marlin_w4a16", + "marlin_w4_weight_repack", + "marlin_w8a16", + "marlin_w8_weight_repack", +] + + +def marlin_w4a16( + inputs: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + zeros: torch.Tensor, + bias: torch.Tensor = None, # TODO + group_size: int = -1, + format: str = "k16n32", + batch_first: bool = True, + outputs: torch.Tensor = None, +): + """ + Args: + inputs: (batch, m, k) if batch_first else (m, batch, k) torch.float16, torch.bfloat16 + weights: (batch, k/16, n/32, 64) torch.int32 + scales: + (batch, k_groups, n) format:k16n32 torch.float16, torch.bfloat16 + (batch, n_groups, k) format:k16n32_grouped_n torch.float16, torch.bfloat16 + zeros: + (batch, k_groups, n/8) format:k16n32 torch.int32 + (batch, n_groups, k/8) format:k16n32_grouped_n torch.int32 + group_size: int + group size of quant + format: str + describe format of weight + batch_first: bool + describe format of input and output + Returns: + outputs: (batch, m, n) if batch_first else (m, batch, n) torch.float16, torch.bfloat16 + """ + if outputs is None: + batch, m = ( + (inputs.shape[0], inputs.shape[1]) + if batch_first + else (inputs.shape[1], inputs.shape[0]) + ) + if format.startswith("k16n32"): + n = weights.shape[2] * 32 + outputs = torch.empty( + (batch, m, n) if batch_first else (m, batch, n), + dtype=inputs.dtype, + device=inputs.device, + ) + + ops.infer.marlin_w4a16( + outputs, inputs, weights, scales, zeros, bias, group_size, format, batch_first + ) + return outputs + + +def marlin_w4_weight_repack( + weights: torch.Tensor, + scales: torch.Tensor = None, + zeros: torch.Tensor = None, + weight_format: str = "gptq", + reformat: str = "k16n32", + pack_order: str = "default", + repack_weight: torch.Tensor = None, +): + """ + Args: + weights: + (batch, k, n/8) weight_format:awq torch.int32 + (batch, k/8, n) weight_format:gptq torch.int32 + scales: + (batch, k_groups, n) format:k16n32 torch.float16, torch.bfloat16 + (batch, n_groups, k) format:k16n32_grouped_n torch.float16, torch.bfloat16 + zeros: + (batch, k_groups, n/8) format:k16n32 torch.int32 + (batch, n_groups, k/8) format:k16n32_grouped_n torch.int32 + weight_format: str + describe format of weight + reformat: str + describe format of repacked weight + pack_order: str + describe pack order on a pack unit + Returns: + repack_weight: (batch, k/16, n/32, 64) torch.int32 + """ + assert weight_format in ["gptq", "gptq_grouped_n", "awq"] + assert reformat in ["k16n32", "k16n32_grouped_n"] + assert pack_order in ["default", "02461357", "01234567"] + + if pack_order == "default": + default_order = { + "gptq": "01234567", + "awq": "02461357", + "gptq_grouped_n": "02461357", + } + pack_order = default_order[weight_format] + + if weight_format.startswith("gptq"): + batch, pack_k, n = weights.shape + k = pack_k * 8 + elif weight_format == "awq": + batch, k, pack_n = weights.shape + n = pack_n * 8 + + repack_scales, repack_zeros = None, None + if reformat.startswith("k16n32"): + if repack_weight is None: + repack_weight = torch.empty( + (batch, k // 16, n // 32, 64), + dtype=torch.int32, + device=weights.device, + ) + if scales is not None: + repack_scales = torch.empty_like(scales) + if zeros is not None: + repack_zeros = torch.empty_like(zeros) + + ops.infer.marlin_w4_weight_repack( + weights, + repack_weight, + scales, + repack_scales, + zeros, + repack_zeros, + weight_format, + reformat, + pack_order, + ) + + if repack_scales is not None and repack_zeros is not None: + return repack_weight, repack_scales, repack_zeros + else: + return repack_weight + + +def marlin_w8a16( + inputs: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + bias: torch.Tensor = None, # TODO + group_size: int = -1, + format: str = "k16n16", + batch_first: bool = True, + outputs: torch.Tensor = None, +): + """ + Args: + inputs: (batch, m, k) if batch_first else (m, batch, k) torch.float16, torch.bfloat16 + weights: (batch, k/16, n/16, 64) torch.int32 + scales: + (batch, k_groups, n) format:k16n16 torch.float32 + (batch, n_groups, k) format:k16n16_grouped_n torch.float32 + group_size: int + group size of quant + format: str + describe format of weight + batch_first: bool + describe format of input and output + Returns: + outputs: (batch, m, n) if batch_first else (m, batch, n) torch.float16, torch.bfloat16 + """ + if outputs is None: + batch, m = ( + (inputs.shape[0], inputs.shape[1]) + if batch_first + else (inputs.shape[1], inputs.shape[0]) + ) + if format.startswith("k16n16"): + n = weights.shape[2] * 16 + outputs = torch.empty( + (batch, m, n) if batch_first else (m, batch, n), + dtype=inputs.dtype, + device=inputs.device, + ) + + ops.infer.marlin_w8a16( + outputs, inputs, weights, scales, bias, group_size, format, batch_first + ) + return outputs + + +def marlin_w8_weight_repack( + weights: torch.Tensor, + scales: torch.Tensor = None, + weight_format: str = "int8", + reformat: str = "k16n16", +): + """ + Args: + weights: + (batch, k, n) weight_format:int8 torch.int8 + scales: + (batch, k_groups, n) format:k16n16 torch.float32 + (batch, n_groups, k) format:k16n16_grouped_n torch.float32 + weight_format: str + describe format of weight + reformat: str + describe format of repacked weight + Returns: + repack_weight: + (batch, k/16, n/16, 64) torch.int32 + """ + assert weight_format in ["int8"] + assert reformat in ["k16n16", "k16n16_grouped_n"] + + repack_scales = None + if weight_format == "int8": + batch, k, n = weights.shape + repack_weight = torch.empty( + (batch, k // 16, n // 16, 64), + dtype=torch.int32, + device=weights.device, + ) + if scales is not None: + repack_scales = torch.empty_like(scales) + + ops.infer.marlin_w8_weight_repack( + weights, + repack_weight, + scales, + repack_scales, + weight_format, + reformat, + ) + + if repack_scales is not None: + return repack_weight, repack_scales + else: + return repack_weight diff --git a/ixformer_sdk/inference/functions/matmul.py b/ixformer_sdk/inference/functions/matmul.py new file mode 100644 index 00000000..0d9c4bd3 --- /dev/null +++ b/ixformer_sdk/inference/functions/matmul.py @@ -0,0 +1,50 @@ +import ixformer._C as ops +import torch + +__all__ = ["matmul", "ref_matmul"] + + +def ref_matmul(input, other, *, transa, transb, alpha): + if transa: + dims = list(range(input.ndim)) + dims[-1], dims[-2] = dims[-2], dims[-1] + input = input.permute(*dims).contiguous() + + if transb: + dims = list(range(other.ndim)) + dims[-1], dims[-2] = dims[-2], dims[-1] + other = other.permute(*dims).contiguous() + + return alpha * torch.matmul(input, other) + + +def matmul( + input: torch.Tensor, + other: torch.Tensor, + *, + transa: bool = False, + transb: bool = False, + alpha: float = 1.0, +) -> torch.Tensor: + """ + Args: + input: (...,m,k) or (...,k,m) torch.half + 当transa为False shape : [...,m,k], 当transa为True shape : [...,k,m] + other: (...,k,n) or (...,n,k) torch.half + 当transa为False shape : [...,m,k], 当transa为True shape : [...,k,m] + transa: bool + transb: bool + alpha: float + Returns: + Tensor: (..., m, n) torch.half + """ + if not input.is_contiguous(): + input = input.contiguous() + + if not other.is_contiguous(): + if not other.transpose(-2, -1).is_contiguous(): + other = other.contiguous() + + return ops.train.matmul( + input, other, transa=transa, transb=transb, alpha=alpha, beta=0.0 + ) diff --git a/ixformer_sdk/inference/functions/mla_fused.py b/ixformer_sdk/inference/functions/mla_fused.py new file mode 100644 index 00000000..4a04901c --- /dev/null +++ b/ixformer_sdk/inference/functions/mla_fused.py @@ -0,0 +1,325 @@ +from typing import Optional + +import ixformer._C as ops +import torch + +__all__ = [ + # 0.6.3 + "ref_minicpm3_fused_rope", + "ref_minicpm3_fused_copy_kv", + "minicpm3_fused_rope", + "minicpm3_fused_copy_kv", + # 0.6.6 + "ref_mla_rope_phi", + "mla_rope_phi", + "ref_mla_rope", + "mla_rope", + "ref_mla_copy_kv", + "mla_copy_kv", +] + + +def _rotate_neox(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _rotate_gptj(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., ::2] + x2 = x[..., 1::2] + x = torch.stack((-x2, x1), dim=-1) + return x.flatten(-2) + + +# vllm 0.6.3 +def ref_minicpm3_fused_rope( + positions: torch.Tensor, + long_prompt_offset: torch.Tensor, + long_short_cos_sin_cache: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + out_query: Optional[torch.Tensor] = None, + out_key: Optional[torch.Tensor] = None, +): + idx = torch.add(positions, long_prompt_offset) + cos_sin = torch.index_select(long_short_cos_sin_cache, 0, idx) + + cos, sin = cos_sin.chunk(2, dim=-1) + cos = cos.repeat(1, 2).unsqueeze(-2) + sin = sin.repeat(1, 2).unsqueeze(-2) + + out_query = query * cos + _rotate_neox(query) * sin + out_key = key * cos + _rotate_neox(key) * sin + + return out_query, out_key + + +def minicpm3_fused_rope( + positions: torch.Tensor, + long_prompt_offset: torch.Tensor, + long_short_cos_sin_cache: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + out_query: Optional[torch.Tensor] = None, + out_key: Optional[torch.Tensor] = None, +): + """ + Args: + positions: (num_tokens,) torch.int64 + long_prompt_offset: (num_tokens,) torch.int64 + long_short_cos_sin_cache: (max_length, head_dim) torch.float16, torch.bfloat16 + query: (num_tokens, num_q_heads, head_dim) same as long_short_cos_sin_cache + key: (num_tokens, num_kv_heads, head_dim) same as long_short_cos_sin_cache + out_query: same as query + out_key: same as key + Returns: + out_query: same as query + out_key: same as key + """ + + if out_query is None: + out_query = torch.empty_like(query) + if out_key is None: + out_key = torch.empty_like(key) + + ops.infer.minicpm3_fused_rope( + positions, + long_prompt_offset, + long_short_cos_sin_cache, + query, + key, + out_query, + out_key, + ) + return out_query, out_key + + +def ref_minicpm3_fused_copy_kv( + k_nope: torch.Tensor, + k_pe: torch.Tensor, + v: torch.Tensor, + new_k: Optional[torch.Tensor] = None, + new_v: Optional[torch.Tensor] = None, +): + num_tokens, num_heads, k_head_dim = k_nope.shape + head_dim = k_pe.shape[-1] + k_head_dim + v_head_dim = v.shape[-1] + + if new_k is None: + new_k = k_nope.new_empty([num_tokens, num_heads, head_dim]) + if new_v is None: + new_v = k_nope.new_empty([num_tokens, num_heads, head_dim]) + + new_k[:, :, :k_head_dim] = k_nope + new_k[:, :, k_head_dim:] = k_pe + new_v[:, :, :v_head_dim] = v + new_v[:, :, v_head_dim:] = 0 + + return new_k.view(num_tokens, -1), new_v.view(num_tokens, -1) + + +def minicpm3_fused_copy_kv( + k_nope: torch.Tensor, + k_pe: torch.Tensor, + v: torch.Tensor, + new_k: Optional[torch.Tensor] = None, + new_v: Optional[torch.Tensor] = None, +): + """ + Args: + k_nope: (num_tokens, num_heads, k_head_dim) torch.float16, torch.bfloat16 + k_pe: (num_tokens, 1, head_dim - k_head_dim) same as k_nope + v: (num_tokens, num_heads, v_head_dim) same as k_nope + new_k: (num_tokens, num_heads, head_dim) same as k_nope + new_v: (num_tokens, num_heads, head_dim) same as k_nope + Returns: + new_k: (num_tokens, num_heads, head_dim) same as k_nope + new_v: (num_tokens, num_heads, head_dim) same as k_nope + """ + + num_tokens, num_heads, k_head_dim = k_nope.shape + head_dim = k_pe.shape[-1] + k_head_dim + + if new_k is None: + new_k = k_nope.new_empty([num_tokens, num_heads * head_dim]) + if new_v is None: + new_v = k_nope.new_empty([num_tokens, num_heads * head_dim]) + + ops.infer.minicpm3_fused_copy_kv(k_nope, k_pe, v, new_k, new_v) + + return new_k, new_v + + +# vllm 0.6.6 +def ref_mla_rope_phi( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + long_short_cos_sin_cache: torch.Tensor, + k: int, + offsets: Optional[torch.Tensor] = None, +): + long_prompt_offset = ( + torch.any(positions > k).float() * torch.full_like(positions, k) + ).long() + idx = ( + torch.add(positions, long_prompt_offset) + if long_prompt_offset is not None + else positions + ) + + idx = torch.add(idx, offsets) if offsets is not None else idx + cos_sin = torch.index_select(long_short_cos_sin_cache, 0, idx) + + cos, sin = cos_sin.chunk(2, dim=-1) + cos = cos.repeat(1, 2).unsqueeze(-2) + sin = sin.repeat(1, 2).unsqueeze(-2) + + query = query * cos + _rotate_neox(query) * sin + key = key * cos + _rotate_neox(key) * sin + + return query, key + + +def mla_rope_phi( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + key_out: torch.Tensor, + long_short_cos_sin_cache: torch.Tensor, + long_offset: torch.Tensor, + k: int, + offsets: Optional[torch.Tensor] = None, +): + """ + Args: + positions: (num_tokens,) torch.int64 + query: (num_tokens, num_q_heads, head_dim) same as long_short_cos_sin_cache + key: (num_tokens, 1, head_dim) same as long_short_cos_sin_cache + key_out: (num_tokens, num_q_heads, head_dim) same as long_short_cos_sin_cache + long_short_cos_sin_cache: (max_length, head_dim) same as long_short_cos_sin_cache + long_offset: (1,) torch.bool + k: int + offsets: (num_tokens,) + Returns: + query: + key_out: + """ + + ops.infer.mla_rope_phi( + positions, + query, + key, + key_out, + long_short_cos_sin_cache, + long_offset, + k, + offsets, + ) + return query, key_out + + +def ref_mla_rope( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + cos_sin_cache: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + rotary_dim: int = None, + is_neox_style: bool = False, +): + """PyTorch-native implementation equivalent to forward().""" + head_size = query.size(-1) + rotary_dim = rotary_dim or head_size + query_rot = query[..., :rotary_dim] + key_rot = key[..., :rotary_dim] + if rotary_dim < head_size: + query_pass = query[..., rotary_dim:] + key_pass = key[..., rotary_dim:] + + cos_sin = cos_sin_cache[ + torch.add(positions, offsets) if offsets is not None else positions + ] + cos, sin = cos_sin.chunk(2, dim=-1) + if is_neox_style: + cos = cos.repeat(1, 1, 2).unsqueeze(-2) + sin = sin.repeat(1, 1, 2).unsqueeze(-2) + else: + cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) + sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) + + rotate_fn = _rotate_neox if is_neox_style else _rotate_gptj + query_rot = query_rot * cos + rotate_fn(query_rot) * sin + key_rot = key_rot * cos + rotate_fn(key_rot) * sin + + if rotary_dim < head_size: + query = torch.cat((query_rot, query_pass), dim=-1) + key = torch.cat((key_rot, key_pass), dim=-1) + else: + query = query_rot + key = key_rot + return query, key + + +def mla_rope( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + key_out: torch.Tensor, + cos_sin_cache: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + is_neox_style: bool = False, +): + """ + Args: + positions: (num_tokens,) torch.int64 + query: (num_tokens, num_q_heads, head_dim) torch.half torch.bfloat torch.float + key: (num_tokens, 1, head_dim) same as query + key_out: (num_tokens, num_q_heads, head_dim) same as query + cos_sin_cache: (max_length, head_dim) same as query + offsets: (num_tokens,) same as query + is_neox_style: bool + Returns: + query: + key_out: + """ + + ops.infer.mla_rope( + positions, + query, + key, + key_out, + cos_sin_cache, + is_neox_style, + offsets, + ) + return query, key_out + + +def ref_mla_copy_kv(key_pe, key_nope, value_nope): + shape = key_nope.shape[:-1] + (key_pe.shape[-1] + key_nope.shape[-1],) + key = torch.empty(shape, device=key_nope.device, dtype=key_nope.dtype) + value = torch.empty_like(key) + + key[..., : key_nope.size(-1)] = key_nope + key[..., key_nope.size(-1) :] = key_pe + value[..., : value_nope.size(-1)] = value_nope + value[..., value_nope.size(-1) :] = 0.0 + return key, value + + +def mla_copy_kv(key_nope, value_nope, key, value): + """ + Args: + key_nope: (num_tokens, num_heads, k_nope_dim) torch.float16, torch.bfloat16, torch.float + value_nope: (num_tokens, num_heads, v_head_dim) same as key_nope + key: (num_tokens, num_heads, head_dim) same as key_nope + value: (num_tokens, num_heads, head_dim) same as key_nope + Returns: + key: + value: + """ + + ops.infer.mla_copy_kv(key_nope, value_nope, key, value) + return key, value diff --git a/ixformer_sdk/inference/functions/mm.py b/ixformer_sdk/inference/functions/mm.py new file mode 100644 index 00000000..eed5e6c5 --- /dev/null +++ b/ixformer_sdk/inference/functions/mm.py @@ -0,0 +1,315 @@ +import ixformer._C as ops +import torch +import torch.nn.functional + +__all__ = [ + "mm", + "addmm", + "fused_addmm_bias_col_act", + "ref_fused_addmm_bias_col_act", + "ref_addmm", + "ref_mm", + "ref_bmm", + "bmm", +] + + +def ref_mm(input, mat, *, out=None): + out = torch.mm(input, mat, out = out) + return out + + +def mm(input, mat, *, out=None): + + """ + Args: + input: (m,k) torch.float16, torch.bfloat16, torch.float32 + mat: (k,n) torch.float16, torch.bfloat16, torch.float32 + out: (m,n) torch.float16, torch.bfloat16, torch.float32 + Returns: + out: (m,n) torch.float16, torch.bfloat16, torch.float32 + """ + assert input.dim() == mat.dim(), "mm tensors must be 2-D" + assert input.size(1) == mat.size( + 0 + ), f"mm cannot be multiplied, {input.size(0)}X{input.size(1)} and {mat.size(0)}X{mat.size(1)}" + + m = input.shape[0] + n = mat.shape[-1] + if out is None: + out = input.new_empty([m, n]) + ops.infer.mm(input, mat, out) + return out + + +"""ixinfer support activations +/// @ingroup GEMM +typedef enum { + CUINFER_BLAS_GEMM_CUSTOM_NONE = 0, + CUINFER_BLAS_GEMM_CUSTOM_BIAS_ADD_ROW_OUT = 1, + CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS = 2, + CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_GELU = 3, + CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_RELU = 4, + CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_TRANSPOSE = 5, + CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS = 6, + CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_GELU = 7, + CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_RELU = 8, + CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_TRANSPOSE = 9, + CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_SIGMOID = 10, + CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_SIGMOID = 11, + CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_SILU = 12, + CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_SILU = 13, + CUINFER_BLAS_GEMM_CUSTOM_SIGMOID = 14, + CUINFER_BLAS_GEMM_CUSTOM_SILU = 15, + CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_TANH = 16, + CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_TANH = 17, + CUINFER_BLAS_GEMM_SPECIAL_INT8_FLOATBIAS = 18, + CUINFER_BLAS_GEMM_SPECIAL_INT8_FLOATBIAS_GELU = 19, + CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_SWISH = 20, + CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_ERF_GELU = 21 +} cuinferGEMMCustomOption_t; +""" + +activation_to_id = { + "fused_bias_col": 2, # support bf16, fp16 + "fused_bias_gelu": 3, # support fp16 + "fused_bias_relu": 4, # support fp16 +} + +id_to_activation = {value: key for key, value in activation_to_id.items()} + + +def ref_addmm(input, mat1, mat2, *, beta=1, alpha=1, out=None): + output_pt = torch.addmm(input, mat1, mat2, alpha=alpha, beta=beta) + return output_pt + + +def ref_fused_addmm_bias_col_act( + input, mat1, mat2, *, beta=1, alpha=1, out=None, bias=None, activation=2 +): + if isinstance(activation, int): + assert activation in id_to_activation + if isinstance(activation, str): + assert activation in activation_to_id + activation = activation_to_id[activation] + if activation == 2: + output_pt = torch.addmm(input, mat1, mat2, alpha=alpha, beta=beta) + bias + elif activation == 3: + output_pt = torch.nn.functional.gelu( + torch.addmm(input, mat1, mat2, alpha=alpha, beta=beta) + bias + ) + else: + output_pt = torch.nn.functional.relu( + torch.addmm(input, mat1, mat2, alpha=alpha, beta=beta) + bias + ) + return output_pt + + +def addmm(input, mat1, mat2, *, beta=1, alpha=1, out=None): + + """ + Args: + input: (m,n) torch.float16, torch.bfloat16, torch.float32 + mat1: (m,k) torch.float16, torch.bfloat16, torch.float32 + mat2: (k,n) torch.float16, torch.bfloat16, torch.float32 + beta: float + alpha: float + out: (m,n) torch.float16, torch.bfloat16, torch.float32 + Returns: + out: (m,n) torch.float16, torch.bfloat16, torch.float32 + """ + assert mat1.dim() == mat2.dim(), "addmm mat1 mat2 tensors must be 2-D" + assert mat1.size(1) == mat2.size( + 0 + ), f"addmm cannot be multiplied, {mat1.size(0)}X{mat1.size(1)} and {mat2.size(0)}X{mat2.size(1)}" + + m = mat1.shape[0] + n = mat2.shape[-1] + + if input is not None and len(input.shape) == 1: + input = input.view(1, -1) + + if out is None: + out = input.new_empty([m, n]) + if input is None: + input = out + beta = 0 + + ops.infer.addmm(input, mat1, mat2, beta, alpha, out, None, 0) + return out + + +def fused_addmm_bias_col_act( + input, mat1, mat2, *, beta=1, alpha=1, out=None, bias=None, activation=2 +): + """ + Args: + input: (m,n) torch.float16, torch.bfloat16, torch.float32 + mat1: (m,k) torch.float16, torch.bfloat16, torch.float32 + mat2: (k,n) torch.float16, torch.bfloat16, torch.float32 + beta: float + alpha: float + out: (m,n) torch.float16, torch.bfloat16, torch.float32 + 当out的shape为(m,n)时,如果out是is_continouns,则bias 必须为(1,n), 否则,bias为(m,1) + bias: (1,n) or (m,1) + activation: str or int + "fused_bias_col": 2, "fused_bias_gelu": 3, "fused_bias_relu": 4 + Returns: + out: (m,n) torch.float16, torch.bfloat16, torch.float32 + """ + assert mat1.dim() == mat2.dim(), "addmm mat1 mat2 tensors must be 2-D" + assert mat1.size(1) == mat2.size( + 0 + ), f"addmm cannot be multiplied, {mat1.size(0)}X{mat1.size(1)} and {mat2.size(0)}X{mat2.size(1)}" + + if isinstance(activation, int): + assert activation in id_to_activation + if isinstance(activation, str): + assert activation in activation_to_id + activation = activation_to_id[activation] + + m = mat1.shape[0] + n = mat2.shape[-1] + + if out is None: + out = input.new_empty([m, n]) + if input is None: + input = out + beta = 0 + + # activations + if activation == 2: + assert bias is not None + if mat1.dtype == torch.float: + ops.infer.addmm(input, mat1, mat2, beta, alpha, out, None, 0) + out.add_(bias) + return out + elif activation == 3: + assert bias is not None + if mat1.dtype == torch.float: + ops.infer.addmm(input, mat1, mat2, beta, alpha, out, None, 0) + out.add_(bias) + out.copy_(torch.nn.functional.gelu(out)) + return out + elif mat1.dtype == torch.bfloat16: + ops.infer.addmm(input, mat1, mat2, beta, alpha, out, bias, 2) + out.copy_(torch.nn.functional.gelu(out)) + return out + elif activation == 4: + assert bias is not None + if mat1.dtype == torch.float: + ops.infer.addmm(input, mat1, mat2, beta, alpha, out, None, 0) + out.add_(bias) + out.copy_(torch.nn.functional.relu(out)) + return out + elif mat1.dtype == torch.bfloat16: + ops.infer.addmm(input, mat1, mat2, beta, alpha, out, bias, 2) + out.copy_(torch.nn.functional.relu(out)) + return out + + ops.infer.addmm(input, mat1, mat2, beta, alpha, out, bias, activation) + return out + + +def ref_bmm( + input: torch.Tensor, + mat2: torch.Tensor, + alpha: float = 1, + format: str = "NN", + input_scales: torch.Tensor = None, + mat2_scales: torch.Tensor = None, + out_dtype: torch.dtype = None, + out: torch.Tensor = None, +): + if format[1] == "T": + input = input.transpose(-1, -2) + if format[0] == "T": + mat2 = mat2.transpose(-1, -2) + + m = input.size(-2) + n = mat2.size(-1) + + bs = input.size(0) + if input.dtype != torch.int8: + out_dtype = input.dtype + if out is None: + out = torch.empty([bs, m, n], dtype=out_dtype, device=input.device) + + if input.dtype != torch.int8: + torch.bmm(input, mat2, out=out) + if alpha != 1: + out = out * alpha + else: + input = input.float() * input_scales.view(1, -1, 1) + mat2 = mat2.float() * mat2_scales.view(1, 1, -1) + out = torch.bmm(input.float(), mat2.float()) * alpha + out = out.to(out_dtype) + return out + + +def bmm( + input: torch.Tensor, + mat2: torch.Tensor, + alpha: float = 1, + format: str = "NN", + input_scales: torch.Tensor = None, + mat2_scales: torch.Tensor = None, + out_dtype: torch.dtype = None, + out: torch.Tensor = None, +): + """ + out = (input@mat2)*alpha + Support three formats: + format: "NN" input shape: (b, m, k) mat2 shape: (b, k, n) out shape: (b, m, n). + If the dtype of input is int8, the following conditions need to be met: n%64==0 k%64==0 + format: "TN" input shape: (b, m, k) mat2 shape: (b, n, k) out shape: (b, m, n) + If the dtype of input is int8, the following conditions need to be met: n%2==0 k%64==0 + format: "NT" input shape: (b, k, m) mat2 shape: (b, k, n) out shape: (b, m, n) + If the dtype of input is int8, the following conditions need to be met: m%64==0 n%64==0 k%64==0 + If the dtype of input is int8, it is necessary to specify out_dtype. + Args: + input: (b, m, k) or (b, k, m) torch.float16, torch.bfloat16, int8 + mat2: (b, k, n) or (b, n, k) torch.float16, torch.bfloat16, int8 + alpha: float32 + format: TN,NN,NT string + input_scales: (m) torch.float32 + mat2_scales: (n) torch.float32 + out_dtype: torch.float16, torch.bfloat16 + out: (b, m, n) torch.float16, torch.bfloat16 + Returns: + out: (b, m, n) torch.float16, torch.bfloat16 + """ + + if format[1] == "N": + m = input.size(-2) + k = input.size(-1) + else: + m = input.size(-1) + k = input.size(-2) + if format[0] == "N": + n = mat2.size(-1) + else: + n = mat2.size(-2) + + if input.dtype == torch.int8: + if format == "TN": + assert ( + n % 2 == 0 and k % 64 == 0 + ), f"bmm shape error, m={m} n={n} k={k}." + elif format == "NT": + assert ( + m % 64 == 0 and n % 64 == 0 and k % 64 == 0 + ), f"bmm shape error, m={m} n={n} k={k}." + elif format == "NN": + assert ( + n % 64 == 0 and k % 64 == 0 + ), f"bmm shape error, m={m} n={n} k={k}." + bs = input.size(0) + if out is None: + if input.dtype != torch.int8: + out_dtype = input.dtype + else: + assert out_dtype is not None + out = torch.empty([bs, m, n], dtype=out_dtype, device=input.device) + ops.infer.bmm(input, mat2, input_scales, mat2_scales, alpha, format, out) + return out diff --git a/ixformer_sdk/inference/functions/moe.py b/ixformer_sdk/inference/functions/moe.py new file mode 100644 index 00000000..80e045e7 --- /dev/null +++ b/ixformer_sdk/inference/functions/moe.py @@ -0,0 +1,1380 @@ +import os +from typing import Optional, Tuple + +import ixformer._C as ops +import torch + +__all__ = [ + "ref_moe_output_reduce_sum", + "moe_output_reduce_sum", + "moe_expand_input", + "ref_moe_expand_input", + "moe_expand_input_dynamic_scaled_int8", + "ref_moe_expand_input_dynamic_scaled_int8", + "moe_compute_token_index", + "moe_compute_token_index_ep", + "ref_moe_compute_token_index_ep", + "moe_topk_softmax", + "ref_moe_topk_softmax", + "moe_grouped_topk", + "ref_moe_grouped_topk", + "moe_align_token_index", + "ref_moe_align_token_index", + "ref_activation_dynamic_scaled_int8", + "activation_dynamic_scaled_int8", + "moe_w8a8_group_gemm", + "ref_moe_w8a8_group_gemm", + "moe_w4a8_group_gemm", + "moe_w4a8_group_gemv", + "ref_moe_w4a8_group_gemm", + "quant_repack_int4", + "moe_w4a16_group_gemm", + "ref_moe_w4a16_group_gemm", +] + + +def ref_moe_output_reduce_sum( + input: torch.Tensor, + topk_weight: torch.Tensor = None, + output: torch.Tensor = None, + mask: torch.Tensor = None, + extra_residual: torch.Tensor = None, + scaling_factor: float = 1.0, +): + if output is None: + m, topk, k = input.shape + output = torch.empty([m, k], dtype=input.dtype, device=input.device) + temp = input.clone().to(torch.float32) + if topk_weight is not None: + temp *= topk_weight.unsqueeze(-1) + if mask is not None: + mask = mask.reshape(m, topk) + mask_value = torch.where(mask, 0.0, 1.0) + temp *= mask_value.unsqueeze(-1) + + temp = torch.sum(temp, dim=1) + if extra_residual is not None: + temp = temp * scaling_factor + extra_residual + output.copy_(temp.to(input.dtype)) + return output + + +def moe_output_reduce_sum( + input: torch.Tensor, + topk_weight: torch.Tensor = None, + output: torch.Tensor = None, + mask: torch.Tensor = None, + extra_residual: torch.Tensor = None, + scaling_factor: float = 1.0, +): + """ + Args: + input: (m, topk, k) torch.float16, torch.bfloat16 + topk_weight: (m, topk) torch.float32 + mask: (m * topk) torch.bool + extra_residual: (m, k) torch.float16, torch.bfloat16 + scaling_factor: float32 + scaling factor for output + Returns: + output: (m, k) torch.float16, torch.bfloat16 + """ + if output is None: + m, topk, k = input.shape + output = torch.empty([m, k], dtype=input.dtype, device=input.device) + ops.infer.moe_output_reduce_sum( + output, input, topk_weight, mask, extra_residual, scaling_factor + ) + return output + + +def ref_moe_expand_input( + hidden_states: torch.Tensor, + dst_to_src: torch.Tensor, + dst_tokens: int, + topk: int, + src_to_dst: torch.Tensor = None, + output: torch.Tensor = None, +): + src_tokens, hidden_size = hidden_states.shape + input_expand = ( + hidden_states.view(src_tokens, 1, hidden_size) + .repeat(1, topk, 1) + .reshape(-1, hidden_size) + ) + if output is None: + output = input_expand[dst_to_src] + else: + output.copy_(input_expand[dst_to_src]) + return output + + +def moe_expand_input( + hidden_states: torch.Tensor, + dst_to_src: torch.Tensor, + dst_tokens: int, + topk: int, + src_to_dst: torch.Tensor = None, + output: torch.Tensor = None, +): + """ + Args: + hidden_states: (num_tokens, hidden_size) torch.float16, torch.bfloat16 + dst_to_src: (num_tokens*topk) torch.int32 + index of dst to src. + dst_tokens: int + the number of tokens after expansion. + topk: int + topk for moe + src_to_dst: (num_tokens*topk) torch.int32 + index of src to dst. + Returns: + output: (dst_tokens, hidden_size) hidden_states.dtype + """ + src_tokens, hidden_size = hidden_states.shape + if output is None: + output = torch.empty( + (dst_tokens, hidden_size), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + ops.infer.moe_expand_input( + output, + hidden_states, + dst_to_src, + src_to_dst, + dst_tokens, + topk, + ) + return output + + +def ref_moe_expand_input_dynamic_scaled_int8( + hidden_states: torch.Tensor, + dst_to_src: torch.Tensor, + dst_tokens: int, + topk: int, + src_to_dst: torch.Tensor = None, + topk_ids: torch.Tensor = None, + smooth_scales: torch.Tensor = None, + i8_output: torch.Tensor = None, + output_scales: torch.Tensor = None, +): + import ixformer.functions as F + + src_tokens, hidden_size = hidden_states.shape + expand_tokens = src_tokens * topk + input_expand = ( + hidden_states.view(src_tokens, 1, hidden_size) + .repeat(1, topk, 1) + .reshape(-1, hidden_size) + ) + + if smooth_scales is not None and topk_ids is not None: + input_expand = input_expand * smooth_scales[topk_ids.flatten()] + input_expand = input_expand.to(hidden_states.dtype) + + intput_i8 = torch.empty( + (expand_tokens, hidden_size), dtype=torch.int8, device=hidden_states.device + ) + input_scales = torch.empty( + expand_tokens, dtype=torch.float32, device=hidden_states.device + ) + F.dynamic_scaled_int8_quant(intput_i8, input_expand, input_scales) + + if i8_output is None: + i8_output = torch.zeros( + (dst_tokens, hidden_size), dtype=torch.int8, device=hidden_states.device + ) + if output_scales is None: + output_scales = torch.zeros( + dst_tokens, dtype=torch.float32, device=hidden_states.device + ) + + i8_output = intput_i8[dst_to_src] + output_scales = input_scales[dst_to_src] + + return i8_output, output_scales + + +def moe_expand_input_dynamic_scaled_int8( + hidden_states: torch.Tensor, + dst_to_src: torch.Tensor, + dst_tokens: int, + topk: int, + src_to_dst: torch.Tensor = None, + topk_ids: torch.Tensor = None, + smooth_scales: torch.Tensor = None, + output_format: int = 0, + i8_output: torch.Tensor = None, + output_scales: torch.Tensor = None, +): + """ + Args: + hidden_states: (num_tokens, hidden_size) torch.float16, torch.bfloat16 + dst_to_src: (num_tokens*topk) torch.int32 + index of dst to src. + dst_tokens: int + the number of tokens after expansion. + topk: int + topk for moe + src_to_dst: (num_tokens*topk) torch.int32 + index of src to dst. + topk_ids: (num_tokens, topk) torch.int32 + smooth_scales: (num_experts, hidden_size) torch.float16, torch.bfloat16 + output_format: int + specific output format for subsequent kernel + 0 : origin output + 1 : used for w4a8 group gemv + Returns: + i8_output: (dst_tokens, hidden_size) torch.int8 + output_scales: (dst_tokens) torch.float32 + """ + hidden_size = hidden_states.shape[-1] + if i8_output is None: + i8_output = torch.empty( + (dst_tokens, hidden_size), dtype=torch.int8, device=hidden_states.device + ) + if output_scales is None: + output_scales = torch.empty( + dst_tokens, dtype=torch.float32, device=hidden_states.device + ) + ops.infer.moe_expand_input_dynamic_scaled_int8( + i8_output, + output_scales, + hidden_states.view(-1, hidden_size), + dst_to_src, + src_to_dst, + topk_ids, + smooth_scales, + dst_tokens, + topk, + output_format, + ) + + return i8_output, output_scales + + +def moe_compute_token_index( + topk_ids: torch.Tensor, + num_experts: int, + src_dst: torch.Tensor = None, + dst_src: torch.Tensor = None, + expert_sizes_gpu: torch.Tensor = None, + expert_sizes_cpu: torch.Tensor = None, +): + """ + Args: + topk_ids: (num_tokens, topk) torch.int32 + num_experts: int + Returns: + src_dst: (num_tokens*topk) torch.int32 + index of src to dst, e.g. src_tensor[i] = dst_tensor[src_dst[i]] + dst_src: (num_tokens*topk) torch.int32 + index of dst to src. + expert_sizes_gpu: (num_experts) torch.int32 + the number of tokens allocated to each expert. (num_experts) + expert_sizes_cpu: (num_experts) torch.int32 + the number of tokens allocated to each expert. (num_experts) + """ + if src_dst is None: + src_dst = topk_ids.new_empty([topk_ids.numel()]) + if dst_src is None: + dst_src = torch.empty_like(src_dst) + if expert_sizes_gpu is None: + expert_sizes_gpu = topk_ids.new_empty([num_experts]) + + ops.infer.moe_compute_token_index( + topk_ids, + src_dst, + dst_src, + expert_sizes_gpu, + expert_sizes_cpu, + None, + 0, + num_experts, + num_experts, + ) + + return src_dst, dst_src, expert_sizes_gpu, expert_sizes_cpu + + +def ref_moe_topk_softmax( + gating_output: torch.Tensor, + topk: int, + topk_weight: torch.Tensor = None, + topk_ids: torch.Tensor = None, + renormalize: bool = True, +): + score = torch.softmax(gating_output, dim=-1, dtype=torch.float32) + topk_weight, topk_ids = torch.topk(score, topk) + if renormalize: + topk_weight = topk_weight / topk_weight.sum(dim=-1, keepdim=True) + return topk_weight, topk_ids.int() + + +def moe_topk_softmax( + gating_output: torch.Tensor, + topk: int, + topk_weight: torch.Tensor = None, + topk_ids: torch.Tensor = None, + renormalize: bool = True, +): + """ + Args: + gating_output: (num_tokens, num_experts) torch.float32 + topk: int + renormalize: bool + Returns: + topk_weight: (num_tokens, topk) torch.float32 + topk_ids: (num_tokens, topk) torch.int32 + """ + num_tokens, num_experts = gating_output.shape + device = gating_output.device + if topk_weight is None: + topk_weight = torch.empty( + [num_tokens, topk], dtype=torch.float32, device=device + ) + if topk_ids is None: + topk_ids = torch.empty([num_tokens, topk], dtype=torch.int32, device=device) + token_expert_indicies = torch.empty( + [num_tokens, topk], dtype=torch.int32, device="cuda" + ) # not use + + ops.infer.moe_topk_softmax( + topk_weight, topk_ids, token_expert_indicies, gating_output, renormalize + ) + return topk_weight, topk_ids + + +def ref_moe_grouped_topk( + gating_output: torch.Tensor, + topk: int, + num_expert_group: int = 0, + topk_group: int = 0, + scoring_func: str = "softmax", + e_score_correction_bias: Optional[torch.Tensor] = None, + renormalize: bool = True, +): + + gating_output = gating_output.to(torch.float32) + if scoring_func == "softmax": + scores = torch.softmax(gating_output, dim=-1) + elif scoring_func == "sigmoid": + scores = gating_output.sigmoid() + else: + raise ValueError(f"Unsupported scoring function: {scoring_func}") + + if e_score_correction_bias is not None: + original_scores = scores + scores = scores + e_score_correction_bias.unsqueeze(0) + + num_token = scores.shape[0] + group_scores = ( + scores.view(num_token, num_expert_group, -1).max(dim=-1).values + ) # [n, n_group] + + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[ + 1 + ] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) + .reshape(num_token, -1) + ) # [n, e] + + tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e] + + if e_score_correction_bias is not None: + topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1] + # Use original unbiased scores for the routing weights + topk_weights = original_scores.gather(1, topk_ids) + else: + topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False) + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + return topk_weights.to(torch.float32), topk_ids.to(torch.int32) + + +def moe_grouped_topk( + gating_output: torch.Tensor, + topk: int, + num_expert_group: int, + topk_group: int, + scoring_func: str = "softmax", + e_score_correction_bias: torch.Tensor = None, + topk_weight: torch.Tensor = None, + topk_ids: torch.Tensor = None, + renormalize: bool = True, +): + """ + Args: + gating_output: (num_tokens, num_experts) torch.float32, torch.bfloat16 + topk: int + num_expert_group: int + topk_group: int + scoring_func: str + e_score_correction_bias: (num_experts) torch.float16, torch.bfloat16 + renormalize: bool + Returns: + topk_weight: (num_tokens, topk) torch.float32 + topk_ids: (num_tokens, topk) torch.int32 torch.int64 + """ + num_tokens, num_experts = gating_output.shape + device = gating_output.device + if topk_weight is None: + topk_weight = torch.empty( + [num_tokens, topk], dtype=torch.float32, device=device + ) + if topk_ids is None: + topk_ids = torch.empty([num_tokens, topk], dtype=torch.int32, device=device) + + ops.infer.moe_grouped_topk( + topk_weight, + topk_ids, + gating_output, + e_score_correction_bias, + num_expert_group, + topk_group, + scoring_func, + renormalize, + ) + return topk_weight, topk_ids + + +def ref_moe_align_token_index( + topk_ids: torch.Tensor, + num_experts: int, + src_to_dst: torch.Tensor = None, + dst_to_src: torch.Tensor = None, + expert_sizes_gpu: torch.Tensor = None, + expert_sizes_cpu: torch.Tensor = None, +): + src_to_dst = [] + dst_to_src = [-1 for _ in range(topk_ids.numel())] + expert_sizes_gpu = torch.empty([num_experts], dtype=torch.int32, device="cuda") + + for i in range(num_experts): + expert_sizes_gpu[i] = (topk_ids == i).sum() + + expert_sizes_gpu_cu = torch.zeros( + [num_experts + 1], dtype=torch.int32, device="cuda" + ) + expert_sizes_gpu_cu[1:] = expert_sizes_gpu + + expert_sizes_gpu_cu = expert_sizes_gpu_cu.cumsum(dim=-1).cpu().tolist() + + topk_ids = topk_ids.view(-1).cpu().tolist() + for i, expert_id in enumerate(topk_ids): + dst_idx = expert_sizes_gpu_cu[expert_id] + expert_sizes_gpu_cu[expert_id] += 1 + src_to_dst.append(dst_idx) + dst_to_src[dst_idx] = i + + src_to_dst = torch.tensor(src_to_dst, dtype=torch.int32, device="cuda") + dst_to_src = torch.tensor(dst_to_src, dtype=torch.int32, device="cuda") + expert_sizes_cpu = expert_sizes_gpu.cpu() + + return src_to_dst, dst_to_src, expert_sizes_gpu, expert_sizes_cpu + + +def moe_align_token_index( + topk_ids: torch.Tensor, + num_experts: int, + src_to_dst: torch.Tensor = None, + dst_to_src: torch.Tensor = None, + expert_sizes_gpu: torch.Tensor = None, + expert_sizes_cpu: torch.Tensor = None, +): + """ + Args: + topk_ids: (num_tokens, topk) torch.int32 + num_experts: int + Returns: + src_to_dst: (num_tokens*topk) torch.int32 + index of src to dst, e.g. src_tensor[i] = dst_tensor[src_to_dst[i]] + dst_to_src: (num_tokens*topk) torch.int32 + index of dst to src. + expert_sizes_gpu: (num_experts) torch.int32 + the number of tokens allocated to each expert. (num_experts) + expert_sizes_cpu: (num_experts) torch.int32 + the number of tokens allocated to each expert. (num_experts) + """ + if src_to_dst is None: + src_to_dst = topk_ids.new_empty([topk_ids.numel()]) + if dst_to_src is None: + dst_to_src = torch.empty_like(src_to_dst) + if expert_sizes_gpu is None: + expert_sizes_gpu = topk_ids.new_empty([num_experts]) + + ops.infer.moe_compute_token_index( + topk_ids, + src_to_dst, + dst_to_src, + expert_sizes_gpu, + expert_sizes_cpu, + None, + 0, + num_experts, + num_experts, + ) + + if expert_sizes_cpu is None: + expert_sizes_cpu = expert_sizes_gpu.detach().cpu() + + return src_to_dst, dst_to_src, expert_sizes_gpu, expert_sizes_cpu + + +def ref_moe_compute_token_index_ep( + topk_ids: torch.Tensor, + num_experts: int, + start_expert_id: int, + end_expert_id: int, + src_to_dst: torch.Tensor = None, + dst_to_src: torch.Tensor = None, + expert_sizes_gpu: torch.Tensor = None, + expert_sizes_cpu: torch.Tensor = None, +): + vaild_num_experts = end_expert_id - start_expert_id + expert_sizes_gpu = torch.empty( + [vaild_num_experts], dtype=torch.int32, device="cuda" + ) + + for i in range(vaild_num_experts): + expert_sizes_gpu[i] = (topk_ids == (i + start_expert_id)).sum() + + expert_sizes_gpu_cu = torch.zeros( + [vaild_num_experts + 1], dtype=torch.int32, device="cuda" + ) + expert_sizes_gpu_cu[1:] = expert_sizes_gpu + + expert_sizes_gpu_cu = expert_sizes_gpu_cu.cumsum(dim=-1).cpu().tolist() + expand_tokens = expert_sizes_gpu_cu[-1] + topk_ids = topk_ids.view(-1).cpu().tolist() + src_to_dst = [] + dst_to_src = [-1 for _ in range(expand_tokens)] + for i, expert_id in enumerate(topk_ids): + if expert_id >= start_expert_id and expert_id < end_expert_id: + eid = expert_id - start_expert_id + dst_idx = expert_sizes_gpu_cu[eid] + expert_sizes_gpu_cu[eid] += 1 + src_to_dst.append(dst_idx) + dst_to_src[dst_idx] = i + else: + src_to_dst.append(-1) + src_to_dst = torch.tensor(src_to_dst, dtype=torch.int32, device="cuda") + dst_to_src = torch.tensor(dst_to_src, dtype=torch.int32, device="cuda") + expert_sizes_cpu = expert_sizes_gpu.cpu() + + return src_to_dst, dst_to_src, expert_sizes_gpu, expert_sizes_cpu, expand_tokens + + +def moe_compute_token_index_ep( + topk_ids: torch.Tensor, + num_experts: int, + start_expert_id: int, + end_expert_id: int, + src_to_dst: torch.Tensor = None, + dst_to_src: torch.Tensor = None, + expert_sizes_gpu: torch.Tensor = None, + expert_sizes_cpu: torch.Tensor = None, +): + """ + Args: + topk_ids: (num_tokens, topk) torch.int32 + num_experts: int + the number of tokens overall + start_expert_id int + start expert id of the vaild expert interval + end_expert_id int + end expert id of the vaild expert interval + Returns: + src_to_dst: (num_tokens*topk) torch.int32 + index of src to dst, e.g. src_tensor[i] = dst_tensor[src_to_dst[i]] + dst_to_src: (expand_tokens_ep) torch.int32 + index of dst to src. + expert_sizes_gpu: (expand_tokens_ep) torch.int32 + the number of tokens allocated to each expert. (num_experts) + expert_sizes_cpu: (expand_tokens_ep) torch.int32 + the number of tokens allocated to each expert. (num_experts) + expand_tokens the number of tokens which expert id in [start_expert_id, end_expert_id) + int + """ + vaild_num_experts = end_expert_id - start_expert_id + if src_to_dst is None: + src_to_dst = topk_ids.new_empty([topk_ids.numel()]) + if dst_to_src is None: + dst_to_src = torch.empty_like(src_to_dst) + if expert_sizes_gpu is None: + expert_sizes_gpu = topk_ids.new_empty([vaild_num_experts]) + expand_tokens_gpu = torch.empty((1), dtype=torch.int32, device=topk_ids.device) + ops.infer.moe_compute_token_index( + topk_ids, + src_to_dst, + dst_to_src, + expert_sizes_gpu, + expert_sizes_cpu, + expand_tokens_gpu, + start_expert_id, + end_expert_id, + num_experts, + ) + + if expert_sizes_cpu is None: + expert_sizes_cpu = expert_sizes_gpu.detach().cpu() + expand_tokens = expand_tokens_gpu.cpu().item() + + return ( + src_to_dst, + dst_to_src[:expand_tokens], + expert_sizes_gpu, + expert_sizes_cpu, + expand_tokens, + ) + + +def ref_activation_dynamic_scaled_int8( + input: torch.Tensor, + bias: torch.Tensor = None, + smooth_scales: torch.Tensor = None, + dst_to_src: torch.Tensor = None, + topk_ids: torch.Tensor = None, + act_type: str = "silu", + i8_output: torch.Tensor = None, + output_scales: torch.Tensor = None, +): + if i8_output is None: + output_shape = ( + input.shape[:-1] + (input.shape[-1] // 2,) + if act_type == "swiglu" + else input.shape + ) + i8_output = torch.empty(output_shape, dtype=torch.int8, device=input.device) + if output_scales is None: + output_scales = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + + temp = input.clone().to(torch.float32) # m, k + + # Add bias + if bias is not None: + if topk_ids is not None and dst_to_src is not None: + # bias (num_experts, k) + temp += bias[topk_ids.flatten()[dst_to_src]] + else: + # bias (k) + temp += bias.view(1, -1) + + # Activation + if act_type == "silu": + temp = torch.nn.functional.silu(temp) + elif act_type == "gelu": + temp = torch.nn.functional.gelu(temp) + elif act_type == "swiglu": + x1, x2 = temp.chunk(chunks=2, dim=-1) + temp = torch.nn.functional.silu(x1) * x2 + + # Quant + if smooth_scales is not None: + assert len(smooth_scales.shape) <= 2 + # Multi smooth scale + if len(smooth_scales.shape) == 2: + temp *= smooth_scales[topk_ids.flatten()[dst_to_src]] + else: + temp *= smooth_scales.view(1, -1) + + amax_, _ = torch.max(torch.abs(temp), dim=-1) + output_scales.copy_(amax_ / 127.0) + output = temp / output_scales.view(-1, 1) + output = torch.clamp(torch.round(output), -127, 127).to(torch.int8) + i8_output.copy_(output) + + return i8_output, output_scales + + +def activation_dynamic_scaled_int8( + input: torch.Tensor, + bias: torch.Tensor = None, + smooth_scales: torch.Tensor = None, + dst_to_src: torch.Tensor = None, + topk_ids: torch.Tensor = None, + act_type: str = "silu", + output_format: int = 0, + i8_output: torch.Tensor = None, + output_scales: torch.Tensor = None, +): + """ + Args: + input: (m, k) torch.float16, torch.bfloat16 + bias: (num_experts, k) torch.float32 + smooth_scales: (num_experts, k) or (num_experts, k//2) torch.float16, torch.bfloat16 + if act_type==swiglu, shape=(num_experts, k//2) + dst_to_src: (m) torch.int32 + index of dst to src. + topk_ids: (m) torch.int32 + act_type: str activation type. + Options include gelu, silu, and swiglu. + output_format: int + specific output format for subsequent kernel + 0 : origin output + 1 : used for w4a8 group gemv + Returns: + i8_output: (m, k) or (m, k//2) torch.int8 + if act_type==swiglu, shape=(m, k//2) + output_scales: (m) torch.float32 + """ + if i8_output is None: + output_shape = ( + input.shape[:-1] + (input.shape[-1] // 2,) + if act_type == "swiglu" + else input.shape + ) + i8_output = torch.empty(output_shape, dtype=torch.int8, device=input.device) + if output_scales is None: + output_scales = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + ops.infer.activation_dynamic_scaled_int8( + i8_output, + output_scales, + input, + smooth_scales, + dst_to_src, + topk_ids, + act_type, + bias, + output_format, + ) + + return i8_output, output_scales + + +def ref_moe_w8a8_group_gemm( + input: torch.Tensor, + weight: torch.Tensor, + i_scales: torch.Tensor, + w_scales: torch.Tensor, + output_dtype: torch.dtype, + tokens_per_experts: torch.Tensor, + dst_to_src: torch.Tensor = None, + format: str = "TN", + output: torch.Tensor = None, + group_size=-1, +): + def get_align_size(format): + input_format = format[1] + return 1 if input_format == "N" else 64 + + if output is None: + if output_dtype is None: + raise RuntimeError( + "ref_moe_w8a8_group_gemm need output_dtype argument when output is none." + ) + m = tokens_per_experts.sum() + output = torch.empty( + (m, w_scales.shape[1]), + dtype=output_dtype, + device=input.device, + ) + assert format in ["NN", "TN", "TT", "NT"] + prefix = 0 + out_prefix = 0 + align_size = get_align_size(format) + for eid, n in enumerate(tokens_per_experts): + start, end = prefix, prefix + n + out_start, out_end = out_prefix, out_prefix + n + cur_inputs = ( + input[start:end] if format[1] == "N" else input[:, start:end].T.contiguous() + ) + cur_scales_i = i_scales[start:end].view(-1, 1) + cur_weights = weight[eid] if format[0] == "T" else weight[eid].T.contiguous() + cur_scales_w = w_scales[eid].view(1, -1) + input_f32 = cur_inputs.to(torch.float32) + weight_f32 = cur_weights.to(torch.float32) + if group_size != -1: + w_shape = weight_f32.shape + weight_f32 = weight_f32.view(-1, group_size) + weight_f32 = weight_f32 * cur_scales_w.view(-1, 1) + weight_f32 = weight_f32.view(w_shape) + output[out_start:out_end] = ( + torch.nn.functional.linear(input_f32, weight_f32) * cur_scales_i + ) + else: + output[out_start:out_end] = ( + torch.nn.functional.linear(input_f32, weight_f32) + * cur_scales_i + * cur_scales_w + ) + prefix += (n + align_size - 1) // align_size * align_size + out_prefix += n + if dst_to_src is not None: + tmp = output.clone() + tmp[dst_to_src] = output + output[:] = tmp[:] + return output + + +def moe_w8a8_group_gemm( + input: torch.Tensor, + weight: torch.Tensor, + i_scales: torch.Tensor, + w_scales: torch.Tensor, + output_dtype: torch.dtype, + tokens_per_experts: torch.Tensor, + dst_to_src: torch.Tensor = None, + format: str = "TN", + output: torch.Tensor = None, +): + """ + Args: + input: (m, k) if format[1]=="N" else (k, m) torch.int8 + weight: (n_experts, n, k) if format[0]=="T" else (n_experts, k, n) torch.int8 + i_scales: (m) torch.float32 + w_scales: (n_experts, n) torch.float32 + output_dtype: torch.dtype + support torch.float16 or torch.bfloat16 + tokens_per_experts: (n_experts) torch.int32 + dst_to_src: (m) torch.int32 + index of dst to src. + format: str + format of input and weight + Returns: + output: (sum(tokens_per_experts), n) output_dtype + input and i_scales may be padding when NT format + """ + m = tokens_per_experts.sum() + if output is None: + if output_dtype is None: + raise RuntimeError( + "moe_w8a8_group_gemm need output_dtype argument when output is none." + ) + output = torch.empty( + (m, w_scales.shape[1]), + dtype=output_dtype, + device=input.device, + ) + + ops.infer.moe_w8a8_group_gemm( + output, + input, + weight, + i_scales, + w_scales, + tokens_per_experts, + dst_to_src, + format, + 0, + m, + ) + return output + + +def quant_repack_int4(x, group_size, version, format, isAsymQuant: bool = False): + n_experts, n, k = x.shape + if version == 1: + assert not isAsymQuant + + if group_size == -1: + max_x, _ = torch.max(torch.abs(x), dim=-1, keepdim=True) + scales = torch.round(max_x / 7) + scales[scales < 1e-6] = 1 + out = torch.round(x / scales).clamp(-8, 7).to(torch.int8) + else: + x = x.view(n_experts, -1, group_size) + max_x, _ = torch.max(torch.abs(x), dim=-1, keepdim=True) + scales = torch.round(max_x / 7) + scales[scales < 1e-6] = 1 + out = torch.round(x / scales).clamp(-8, 7).to(torch.int8) + + out = out.view(n_experts, n, k) + + if format[0] == "N": + out = out.transpose(-2, -1).contiguous() # NT (num_experts, k , n) + out = out.reshape(n_experts, k // 32, 2, 16, n // 32, 2, 16) + out = out.view(n_experts, k // 32, 2, 16, n // 32, 2, 16) + out = out.permute(0, 1, 5, 3, 4, 2, 6).contiguous().view(n_experts, k, n) + + ## rearange 32 token + shape = out.shape + out = out.view(shape[0], shape[1], shape[-1] // 32, 32) + out_tmp = out.new_empty(shape[0], shape[1], shape[-1] // 32, 16) + for i in range(16): + sign_low_4bit = (out[:, :, :, i] < 0).to(torch.int8) + low_4bit = sign_low_4bit * 8 + (out[:, :, :, i] & 0x07) + high_4bit = out[:, :, :, i + 16] << 4 + out_tmp[:, :, :, i] = high_4bit + low_4bit + out = out_tmp.view(shape[0], shape[1], shape[-1] // 2).contiguous() + + scales = ( + scales.view(n_experts, n, k // group_size).permute(0, 2, 1).contiguous() + if group_size != -1 + else scales.view(n_experts, n) + ) + + return out, scales, None + + if version == 2: + """ + For group_size == -1 (per-channel), the default scale factor is 18 since + 127 / 7 = 18, for quantization with clip, the scale can be set to 16, 17, etc. + the alpha in ixinfer_gemm_helper need to be set to scale / 16.0, and the ixformer + need to be rebuilt. + """ + if group_size == -1: + out = torch.round(x / 18).clamp(-8, 7).to(torch.int8) + else: + x = x.view(n_experts, -1, group_size) + if isAsymQuant: + max_x, _ = torch.max(x, dim=-1, keepdim=True) + min_x, _ = torch.min(x, dim=-1, keepdim=True) + scales = ((max_x.to(torch.float32) - min_x.to(torch.float32)) / 15).to( + torch.int8 + ) + zeros = (-min_x / scales - 8).to( + torch.int8 + ) # weight use int4 not uint4, and zero use int8 + out = (x / scales + zeros).clamp(-8, 7).to(torch.int8) + else: + max_x, _ = torch.max(torch.abs(x), dim=-1, keepdim=True) + scales = torch.round(max_x / 7) + scales[scales < 1e-6] = 1 + scales = scales.to(torch.int8) + out = torch.round(x / scales).clamp(-8, 7).to(torch.int8) + out = out.view(n_experts, n, k).contiguous() + + if format[0] == "N": + out = out.transpose(-2, -1).contiguous() # NT (num_experts, k , n) + out = out.reshape(n_experts, k // 32, 2, 16, n // 32, 2, 16) + out = out.view(n_experts, k // 32, 2, 16, n // 32, 2, 16) + out = out.permute(0, 1, 5, 3, 4, 2, 6).contiguous().view(n_experts, k, n) + + ## rearange 32 token + shape = out.shape + out = out.view(shape[0], shape[1], shape[-1] // 32, 32) + out_tmp = out.new_empty(shape[0], shape[1], shape[-1] // 32, 16) + for i in range(16): + sign_low_4bit = (out[:, :, :, i] < 0).to(torch.int8) + low_4bit = sign_low_4bit * 8 + (out[:, :, :, i] & 0x07) + high_4bit = out[:, :, :, i + 16] << 4 + out_tmp[:, :, :, i] = high_4bit + low_4bit + out = out_tmp.view(shape[0], shape[1], shape[-1] // 2).contiguous() + + if group_size == -1: + return out, None, None + + scales = scales.to(torch.uint8) + scales_4i8pack = scales.clone().to(torch.int32) + for i in range(3): + scales_4i8pack <<= 8 + scales_4i8pack |= scales + scales_4i8pack = ( + scales_4i8pack.view(n_experts, n, k // group_size) + .permute(0, 2, 1) + .contiguous() + ) + + if not isAsymQuant: + return out, scales_4i8pack, None + + zeros = zeros.to(torch.uint8) + zeros_4i8pack = zeros.clone().to(torch.int32) + for i in range(3): + zeros_4i8pack <<= 8 + zeros_4i8pack |= zeros + zeros_4i8pack = ( + zeros_4i8pack.view(n_experts, n, k // group_size) + .permute(0, 2, 1) + .contiguous() + ) + + return out, scales_4i8pack, zeros_4i8pack + + +def _dequant_weight_int8(tensor, i8scales, i8zeros, group_size, version, format): + """ + format == TN + tensor: (num_experts, n, k/2) + scales: (num_experts, n) if group_size == -1 else (num_experts, k // group_size, n) + format == NT or NN + tensor: (num_experts, k, n/2) + scales: (num_experts, n) if group_size == -1 else (num_experts, k // group_size, n) + output tensor is always k-major + """ + dtype = torch.int8 + + left = (tensor & 0xF0) >> 4 + right = tensor & 0x0F + sign_bit = (tensor >> 3) & 1 + right = (right - (sign_bit * 16)).clamp(-8, 7) + left, right = right, left + + shape = list(left.shape) + left = left.reshape( + shape[:-1] + [shape[-1] // 16, 16] + ) # TN (num_experts, n, k/2/16, 16) + right = right.reshape( + shape[:-1] + [shape[-1] // 16, 16] + ) # TN (num_experts, n, k/2/16, 16) + ret = torch.cat((left, right), dim=-1) # TN (num_experts, n, k/2/16, 32) + ret = ret.reshape( + shape[:-1] + [shape[-1] * 2] + ) # TN (num_experts, n, k); NT (num_experts, k, n) + + ## NT 需要再次转换 + if format[0] == "T": + n_experts, n, k = ret.shape + else: + n_experts, k, n = ret.shape + if format[0] == "N": + ret = ret.view(n_experts, k // 32, 2, 16, n // 32, 2, 16) + ret = ret.permute(0, 1, 5, 3, 4, 2, 6).contiguous().view(n_experts, k, n) + ret = ret.transpose(-2, -1).contiguous() # (num_experts, n, k) + + ret_shape = ret.size() + ret = ret.view(-1, group_size) if group_size != -1 else ret.view(-1, ret.shape[-1]) + + if version == 2: + if group_size == -1: + ret *= 18 # same with quant_repack_int4 + else: + scales = i8scales.to(torch.int8).transpose(-2, -1).contiguous().view(-1, 1) + if i8zeros is not None: + zeros = ( + i8zeros.to(torch.int8).transpose(-2, -1).contiguous().view(-1, 1) + ) + ret -= zeros + ret = scales * ret + + ret = ret.reshape(ret_shape).to(dtype) + return ret + + +def ref_moe_w4a8_group_gemm( + input: torch.Tensor, + weight: torch.Tensor, + i_scales: torch.Tensor, + w_scales: torch.Tensor, + output_dtype: torch.dtype, + tokens_per_experts: torch.Tensor, + w_i8scales: torch.Tensor = None, + w_i8zeros: torch.Tensor = None, + dst_to_src: torch.Tensor = None, + format: int = 0, + version: int = 2, + group_size: int = -1, + persistent: int = 0, + output: torch.Tensor = None, +): + assert format in ["NN", "NT", "TN"], f"w4a8 group gemm only support NN, NT, TN" + + weight_i8 = _dequant_weight_int8( + weight, w_i8scales, w_i8zeros, group_size, version, format + ) + + if format[0] == "N": + weight_i8 = weight_i8.permute(0, 2, 1).contiguous() + if version == 1 and group_size != -1: + w_scales = w_scales.permute(0, 2, 1).contiguous() + + output = ref_moe_w8a8_group_gemm( + input, + weight_i8, + i_scales, + w_scales, + output_dtype, + tokens_per_experts, + dst_to_src, + format, + output, + group_size=-1 if version == 2 else group_size, + ) + return output + + +def moe_w4a8_group_gemv( + input: torch.Tensor, + weight: torch.Tensor, + i_scales: torch.Tensor, + w_scales: torch.Tensor, + output_dtype: torch.dtype, + tokens_per_experts: torch.Tensor, + w_i8scales: torch.Tensor = None, + w_i8zeros: torch.Tensor = None, + dst_to_src: torch.Tensor = None, + format: int = 0, + group_size: int = -1, + persistent: int = 0, + output: torch.Tensor = None, +): + """ + Args: + input: (m, k) torch.int8 + weight: (n_experts, n, k//2) if format & 0b10 else (n_experts, k, n//2) torch.int8 + i_scales: (m) torch.float32 + w_scales: (n_experts, n) if group_size == -1 else (n_experts, k // group_size, n) torch.float32 + output_dtype: torch.float16, torch.bfloat16 + tokens_per_experts: (n_experts) torch.int32 + w_i8scales (n_experts, k // group_size, n) if gorup_size != -1 else (n_experts, n) torch.int32 + w_i8zeros (n_experts, k // group_size, n) if gorup_size != -1 else (n_experts, n) torch.int32 + dst_to_src: (sum(tokens_per_experts)) torch.int32 + format: [0(0b00, NN), 1(0b01, NT), 2(0b10, TN), 3(0b11, TT)] int + group_size version1: NN/NT:[-1,256,320,512], TN:[-1,256,512], + version2: NN:[-1,64], NT:[-1], TN:[-1] int + Returns: + output: (m, n) output_dtype + """ + assert format in [2], f"w4a8 group gemv only support 2(TN)" + + # unsupported EP + # outout_m = tokens_per_experts.sum() + + outout_m = input.size(0) + if output is None: + assert output_dtype is not None, print( + "moe_w4a8_group_gemv need output_dtype argument when output is none." + ) + output = torch.empty( + (outout_m, w_scales.shape[-1]), + dtype=output_dtype, + device=input.device, + ) + + ops.infer.moe_w4a8_group_gemv( + output, + input, + weight, + i_scales, + w_scales, + tokens_per_experts, + w_i8scales, + w_i8zeros, + dst_to_src, + format, + group_size, + persistent, + outout_m, + ) + return output + + +def moe_w4a8_group_gemm( + input: torch.Tensor, + weight: torch.Tensor, + i_scales: torch.Tensor, + w_scales: torch.Tensor, + output_dtype: torch.dtype, + tokens_per_experts: torch.Tensor, + w_i8scales: torch.Tensor = None, + w_i8zeros: torch.Tensor = None, + dst_to_src: torch.Tensor = None, + format: int = 0, + version: int = 2, + group_size: int = -1, + persistent: int = 0, + output: torch.Tensor = None, +): + """ + ColumnMajor matrix :(m, k) @ (k, n) -> (m, n) + RowMajof matrix :(k, m) @ (n, k) -> (n, m) + C = A(weight) @ B(input) + Args: + input: (n, k) torch.int8 + weight: (n_experts, m, k//2) if format & 0b10 else (n_experts, k, m//2) torch.int8 + i_scales: (n) torch.float32 + w_scales: (n_experts, m) if group_size == -1 else (n_experts, k // group_size, m) torch.float32 + output_dtype: torch.float16, torch.bfloat16 + tokens_per_experts: (n_experts) torch.int32 + w_i8scales (n_experts, k // group_size, m) if gorup_size != -1 else (n_experts, m) torch.int32 + w_i8zeros (n_experts, k // group_size, m) if gorup_size != -1 else (n_experts, m) torch.int32 + dst_to_src: (sum(tokens_per_experts)) torch.int32 + format: [0(0b00, NN), 1(0b01, NT), 2(0b10, TN), 3(0b11, TT)] int + version: [1, 2] int + group_size version1: NN/NT:[-1,256,320,512], TN:[-1,256,512], + version2: NN:[-1,64], NT:[-1], TN:[-1] int + Returns: + output: (sum(tokens_per_experts), n) output_dtype + """ + + assert format in [0, 1, 2], f"w4a8 group gemm only support 0(NN),1(NT),2(TN)" + + outout_n = tokens_per_experts.sum() + if output is None: + assert output_dtype is not None, print( + "moe_w4a8_group_gemm need output_dtype argument when output is none." + ) + output = torch.empty( + (outout_n, w_scales.shape[-1]), + dtype=output_dtype, + device=input.device, + ) + + ops.infer.moe_w4a8_group_gemm( + output, + input, + weight, + i_scales, + w_scales, + tokens_per_experts, + w_i8scales, + w_i8zeros, + dst_to_src, + format, + version, + group_size, + outout_n, + persistent, + ) + return output + + +def ref_moe_w4a16_group_gemm( + input: torch.Tensor, + weight: torch.Tensor, + w_scales: torch.Tensor, + quant_type: str, + tokens_per_experts: torch.Tensor, + w_zeros: torch.Tensor = None, + group_size: int = -1, + dst_to_src: torch.Tensor = None, + format: str = "NN", + output: torch.Tensor = None, +): + assert quant_type in ["awq"] + assert format in ["NN"] + from .quantized_linear import ref_quantized_weight_dequant + + output_dtype = input.dtype + + def get_align_size(format): + input_format = format[1] + return 1 if input_format == "N" else 64 + + if output is None: + m = tokens_per_experts.sum() + output = torch.empty( + (m, w_scales.shape[2]), + dtype=output_dtype, + device=input.device, + ) + + prefix = 0 + out_prefix = 0 + align_size = get_align_size(format) + for eid, n in enumerate(tokens_per_experts): + if n == 0: + continue + start, end = prefix, prefix + n + out_start, out_end = out_prefix, out_prefix + n + cur_inputs = ( + input[start:end] if format[1] == "N" else input[:, start:end].T.contiguous() + ) + + cur_weights = weight[eid] + cur_scales = w_scales[eid] + cur_zeros = w_zeros[eid] + + input_f32 = cur_inputs.to(torch.float32) + weight_f32 = ref_quantized_weight_dequant( + qweights=cur_weights, + scales=cur_scales, + quant_type="awq", + output_type=torch.float32, + bits=4, + qzeros=cur_zeros, + group_size=group_size, + g_idx=None, + ) + output[out_start:out_end] = torch.matmul(input_f32, weight_f32) + prefix += (n + align_size - 1) // align_size * align_size + out_prefix += n + + if dst_to_src is not None: + tmp = output.clone() + tmp[dst_to_src] = output + output[:] = tmp[:] + return output + + +def moe_w4a16_group_gemm( + input: torch.Tensor, + weight: torch.Tensor, + w_scales: torch.Tensor, + quant_type: str, + tokens_per_experts: torch.Tensor, + w_zeros: torch.Tensor = None, + group_size: int = -1, + dst_to_src: torch.Tensor = None, + format: str = "NN", + output: torch.Tensor = None, + tokens_per_experts_gpu: torch.Tensor = None, +): + """ + Args: + input: (m, k) torch.float16, torch.bfloat16 + weight: + awq:(n_experts, k, n/8) torch.int32 + w_scales: + awq:(n_experts, k/group_size, n) input.dtype + w_zeros: + awq:(n_experts, k/group_size, n/8) torch.int32 + quant_type: str + quant type for weight, support [awq] now + tokens_per_experts: (n_experts) torch.int32 + group_size: int + dst_to_src: (m) torch.int32 + index of dst to src. + format: str + format of input and weight + Returns: + output: (sum(tokens_per_experts), n) output_dtype + input and i_scales may be padding when NT format + """ + assert quant_type in ["awq"] + assert format in ["NN"] + output_dtype = input.dtype + m = tokens_per_experts.sum() + if output is None: + output = torch.empty( + (m, w_scales.shape[-1]), + dtype=output_dtype, + device=input.device, + ) + FLAG = int(os.getenv("ENABLE_MOE_GROUP_GEMV", 1)) + if FLAG == 1 and tokens_per_experts.max() <= 2: + assert ( + tokens_per_experts_gpu is not None + ), "moe group gemv must have tokens_per_experts_gpu!" + ops.infer.moe_group_gemv( + output, + input, + weight, + w_scales, + tokens_per_experts, + tokens_per_experts_gpu, + w_zeros, + dst_to_src, + quant_type, + format, + group_size, + 0, + m, + ) + return output + + ops.infer.moe_w4a16_group_gemm( + output, + input, + weight, + w_scales, + tokens_per_experts, + w_zeros, + dst_to_src, + quant_type, + format, + group_size, + 0, + m, + ) + return output diff --git a/ixformer_sdk/inference/functions/overlap_comm.py b/ixformer_sdk/inference/functions/overlap_comm.py new file mode 100644 index 00000000..5fcbb23a --- /dev/null +++ b/ixformer_sdk/inference/functions/overlap_comm.py @@ -0,0 +1,84 @@ +import itertools +from functools import partial +from typing import Callable, Dict, Iterable, Tuple + +import torch +import torch.distributed as dist + +import ixformer.distributed as ixfd +from ixformer.core.dispatcher import Dispatcher +from ixformer.core.operator_autotuning import ( + OperatorPreBaseRangeAutotuning, + sync_ranks_metric, +) +from ixformer.distributed import overlap_comm +from ixformer.inference.overlap.linear_mlp_overlap_comm import linear_mlp_overlap +from ixformer.distributed.overlap_comm import GemmMethod + +__all__ = ["linear_allreduce_overlap", "linear_mlp_overlap"] + + +class LinearAllReducePreAutotuning(OperatorPreBaseRangeAutotuning, Dispatcher): + def __init__(self, comm_group, *args, **kwargs): + dist_barrier = True + if "dist_barrier" in kwargs: + dist_barrier = kwargs.pop("dist_barrier") + + super().__init__(dist_barrier=dist_barrier, *args, **kwargs) + self._comm_group = comm_group + self._world_size = ixfd.get_group_world_size(comm_group) + + @classmethod + def dispatcher_key(cls, comm_group, *args, **kwargs): + return (comm_group,) + + def operators(self): + chunks = [2, 4] + gemm_algos = [GemmMethod.kCUINFER, GemmMethod.kCUBLAS, GemmMethod.kLIMITED_GEMM] + candidate_ops = [overlap_comm.GemmAllReduceSplitOverlapComm.native_forward] + for num_chunks, algo in itertools.product(chunks, gemm_algos): + candidate_ops.append( + partial( + overlap_comm.linear_allreduce_overlap, + num_chunks=num_chunks, + gemm_method=algo, + ) + ) + + return candidate_ops + + @property + def _gemm_shapes(self): + basic_k = [4096, 6114, 8192] + tp_k = [k // self._world_size for k in basic_k] + basic_k = tp_k + + basic_m = (512, 1024, 2048, 4096, 8192) + + shapes = set(itertools.product(basic_m, basic_k)) + + return shapes + + def get_operator_key(self, input, *args, **kwargs): + ndim = input.ndim + shape = input.shape + + if ndim == 1: + return (1, shape[0]) + elif ndim == 2: + return shape + else: + return (sum(shape[:-1]), shape[-1]) + + def generate_operator_inputs(self) -> Iterable[Tuple[Tuple, Dict]]: + for m, kn in self._gemm_shapes: + input = torch.randn(m, kn, device="cuda", dtype=torch.half) + weight = torch.randn(kn, kn, device="cuda", dtype=torch.half) + yield (input, weight), {} + + def perf_operator_time(self, op: Callable, *args, **kwargs) -> float: + op_time = super().perf_operator_time(op, *args, **kwargs) + return sync_ranks_metric(op_time, group=self._comm_group) + + +linear_allreduce_overlap = overlap_comm.linear_allreduce_overlap diff --git a/ixformer_sdk/inference/functions/paged_attention.py b/ixformer_sdk/inference/functions/paged_attention.py new file mode 100644 index 00000000..98e985e9 --- /dev/null +++ b/ixformer_sdk/inference/functions/paged_attention.py @@ -0,0 +1,143 @@ +from typing import Union + +import ixformer._C as ops +import torch + +__all__ = [ + "paged_attention", + "paged_attention_flashinfer", + "paged_attention_cache_appended", +] +# paged_attention_cache_append + + +def paged_attention_cache_appended( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_format: str = "HND", # STD NHD HND + key_cache_scales: torch.Tensor = None, + value_cache_scales: torch.Tensor = None, +): + if isinstance(key, torch.Tensor): + ops.infer.paged_attention_cache_appended( + key, + value, + key_cache, + value_cache, + slot_mapping, + key.stride(0), + value.stride(0), + key_cache.stride(0), + value_cache.stride(0), + kv_cache_format, + key_cache_scales, + value_cache_scales, + ) + else: + raise NotImplementedError() + + +def paged_attention( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + max_seq_len: int, + alibi_slopes: torch.Tensor = None, + use_sqrt_alibi: bool = False, + key_cache_scales: torch.Tensor = None, + value_cache_scales: torch.Tensor = None, + kv_cache_format: str = "HND", + algo: int = -1, +): + """ + kv_cache_format + STD : k/v format as same as vllm + NHD : k/v format is [block_size, num_kv_heads, head_dim] in one page + HND : k/v format is [num_kv_heads, block_size, head_dim] in one page + algo + -1 : auto chooes algorithm according to kv_cache_format + 0 : use the first algorithm + 1 : use the second algorithm + """ + if isinstance(query, torch.Tensor): + ops.infer.paged_attention( + output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + block_size, + max_seq_len, + use_sqrt_alibi, + alibi_slopes, + key_cache_scales, + value_cache_scales, + kv_cache_format, + algo, + ) + else: + raise NotImplementedError() + + +def paged_attention_flashinfer( + output: torch.Tensor, + query: torch.Tensor, + paged_kv_data, + paged_kv_indptr: torch.Tensor, + paged_kv_indices: torch.Tensor, + paged_kv_last_page_len: torch.Tensor, + scale: float, + max_seq_len: int = -1, + use_sqrt_alibi: bool = False, + alibi_slopes: torch.Tensor = None, + kv_cache_format: str = "HND", + # key_cache_scales: torch.Tensor = None, + # value_cache_scales: torch.Tensor = None, +): + """ + out / query : [num_seqs, num_qo_heads, head_size] + paged_kv_data + Tensor: + NHD [max_num_pages, 2, page_size, num_kv_heads, head_size] + HND [max_num_pages, 2, num_kv_heads, page_size, head_size] + tuple(k_data, v_data) + NHD [max_num_pages, page_size, num_kv_heads, head_size] + HND [max_num_pages, num_kv_heads, page_size, head_size] + paged_kv_indptr int32 : [num_seqs + 1] + paged_kv_indices int32 : [max_num_pages] + paged_kv_last_page_len int32 : [num_seqs] + """ + if isinstance(paged_kv_data, tuple): + k_data, v_data = paged_kv_data + pack_kv_data = (None, k_data, v_data) + else: + pack_kv_data = (paged_kv_data, None, None) + + if isinstance(query, torch.Tensor): + ops.infer.paged_attention_flashinfer( + output, + query, + *pack_kv_data, + paged_kv_indptr, + paged_kv_indices, + paged_kv_last_page_len, + scale, + max_seq_len, + use_sqrt_alibi, + alibi_slopes, + kv_cache_format, + ) + else: + raise NotImplementedError() diff --git a/ixformer_sdk/inference/functions/quantized_linear.py b/ixformer_sdk/inference/functions/quantized_linear.py new file mode 100644 index 00000000..f24d35ee --- /dev/null +++ b/ixformer_sdk/inference/functions/quantized_linear.py @@ -0,0 +1,238 @@ +from typing import Union + +import ixformer._C as ops +import torch +from torch import Tensor + +__all__ = [ + "quantized_linear", + "quantized_weight_dequant", + "ref_quantized_weight_dequant", + "weight_quantize", +] + + +def quantized_linear( + inputs: torch.Tensor, + qweights: torch.Tensor, + scales: torch.Tensor, + quant_type: str, + bits: int, + qzeros: torch.Tensor = None, + bias: torch.Tensor = None, + group_size: int = -1, + g_idx: torch.Tensor = None, + format: str = "unknown", +): + """ + QuantType inputs qweights Scales bits qzeros bias GroupSize Format ApiCall 备注 + awq (bs, ic) bf16/fp16 int32 NN:(ic, oc // 8) TN:(oc, ic // 8) (ic // group_size, oc)fp16/bf16 4/8 int32(ic // group_size, oc // 8) (oc) or None fp16/bf16 32/128 TN/NN vllm & auto-awq + gptq (bs, ic) bf16/fp16 int32 (ic//8, oc) (ic // group_size, oc)fp16/bf16 4 int32(ic // group_size, oc // 8) (oc) or None fp16/bf16 ic/128 \ auto-gptq bs 只支持到8 + fp4 (bs, ic) bf16/fp16 uint8 (oc * ic // 2, 1) (oc * ic // group_size)fp32 4 \ (oc) or None fp16/bf16 64 \ bitsandbytes bs 只支持到8 + nf4 (bs, ic) bf16/fp16 uint8 (oc * ic // 2, 1) (oc * ic // group_size)fp32 4 \ (oc) or None fp16/bf16 64 \ bitsandbytes bs 只支持到8 + int8 (bs, ic) bf16/fp16 int8 TN:(oc, ic) NN:(ic, oc) (1, oc)fp16/bf16 8 \ (oc) or None fp16/bf16 -1 TN/NN vllm & bitsandbytes + + """ + if isinstance(inputs, torch.Tensor) and not inputs.requires_grad: + return ops.infer.quantized_linear( + inputs, + qweights, + scales, + quant_type, + bits, + qzeros, + bias, + group_size, + g_idx, + format, + ) + raise NotImplementedError() + + +def quantized_weight_dequant( + qweights: torch.Tensor, + scales: torch.Tensor, + quant_type: str, + output_type: str, + bits: int, + qzeros: torch.Tensor = None, + group_size: int = -1, + g_idx: torch.Tensor = None, +): + """ + Args: + qweights: (oc, ic//2) or (ic// (32/bits, oc) torch.unint8 or torch.int32 + scales: (oc * ic//g) or (ic // g, oc) torch.float16, torch.bfloat16, torch.float32 + quant_type: str + 可选项:fp4/nf4/gptq/gptq-ex + output_type: str + 可选项:fp16/bf16 + bits: int + 可选项:4/8 + qzeros: (ic//g, oc//(32/bits)) torch.int32 + group_size: int + 可选项:-1/64/128 + g_idx: (ic) torch.int + + Returns: + Tensor: (oc, ic) or (ic, oc) torch.float16, torch.bfloat16 + + quant_type qweights scales qzeros output_type bits group_size g_idx + fp4/nf4 (oc, ic//2) uint8 (oc * ic//g) fp32 fp16/bf16 / 64/128 / + gptq/gptq-ex (ic// (32/bits, oc) int32 (ic // g, oc) fp16/bf16 (ic // g, oc // (32/bits)) int32 fp16/bf16 4/8 -1 (ic) + + """ + if isinstance(qweights, torch.Tensor) and not qweights.requires_grad: + return ops.infer.quantized_weight_dequant( + qweights, scales, quant_type, output_type, bits, qzeros, group_size, g_idx + ) + raise NotImplementedError() + + +def ref_quantized_weight_dequant( + qweights: torch.Tensor, + scales: torch.Tensor, + quant_type: str, + output_type: torch.dtype, + bits: int, + qzeros: torch.Tensor = None, + group_size: int = -1, + g_idx: torch.Tensor = None, + order_map: list = None, +): + assert quant_type in ["awq"] + if quant_type == "awq": + # qweights:(k, n/8) int32 + # scale:(k/group_size, n) f16 + # qzeros:(k/group_size, n/8) int32 + ic, oc = qweights.shape[0], scales.shape[1] + assert bits == 4 + if order_map is None: + order_map = [0, 2, 4, 6, 1, 3, 5, 7] + order_map = torch.Tensor(order_map).to(torch.int32).to(qweights.device) + order_map = order_map.argsort() + + # (1, 8) + wf = ( + torch.tensor(list(range(0, 32, bits)), dtype=torch.int32) + .unsqueeze(0) + .to(qweights.device) + ) + + # unpack qzeros + unpack_zeros = torch.bitwise_right_shift( + torch.unsqueeze(qzeros, 2).expand(-1, -1, 32 // bits), wf.unsqueeze(0) + ).to(torch.int16 if bits == 8 else torch.int8) + unpack_zeros = unpack_zeros[:, :, order_map] + unpack_zeros = torch.bitwise_and(unpack_zeros, (2**bits) - 1) + # groups, 1, n + unpack_zeros = unpack_zeros.reshape(unpack_zeros.shape[0], 1, -1) + + # unpack weights + unpack_weights = torch.bitwise_right_shift( + torch.unsqueeze(qweights, 2).expand(-1, -1, 32 // bits), + wf.unsqueeze(0), + ).to(torch.int16 if bits == 8 else torch.int8) + unpack_weights = unpack_weights[:, :, order_map] + unpack_weights = torch.bitwise_and(unpack_weights, (2**bits) - 1) + # w : groups, group_size, n + unpack_weights = unpack_weights.reshape( + -1, group_size, unpack_weights.shape[1] * unpack_weights.shape[2] + ) + + deq_weights = (unpack_weights - unpack_zeros) * scales.reshape( + -1, 1, scales.shape[-1] + ) + deq_weights = deq_weights.reshape(ic, oc) + return deq_weights.to(output_type) + + +def create_dynamic_map(signed=True, max_exponent_bits=7, total_bits=8): + """ + Creates the dynamic quantiztion map. + + The dynamic data type is made up of a dynamic exponent and + fraction. As the exponent increase from 0 to -7 the number + of bits available for the fraction shrinks. + + This is a generalization of the dynamic type where a certain + number of the bits and be reserved for the linear quantization + region (the fraction). n determines the maximum number of + exponent bits. + + For more details see + (8-Bit Approximations for Parallelism in Deep Learning)[https://arxiv.org/abs/1511.04561] + """ + + data = [] + # these are additional items that come from the case + # where all the exponent bits are zero and no + # indicator bit is present + non_sign_bits = total_bits - (1 if signed else 1) + additional_items = 2 ** (non_sign_bits - max_exponent_bits) - 1 + for i in range(max_exponent_bits): + fraction_items = int( + 2 ** (i + non_sign_bits - max_exponent_bits) + 1 + if signed + else 2 ** (i + non_sign_bits - max_exponent_bits + 1) + 1, + ) + boundaries = torch.linspace(0.1, 1, fraction_items) + means = (boundaries[:-1] + boundaries[1:]) / 2.0 + data += ((10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() + if signed: + data += (-(10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() + + if additional_items > 0: + boundaries = torch.linspace(0.1, 1, additional_items + 1) + means = (boundaries[:-1] + boundaries[1:]) / 2.0 + data += ((10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() + if signed: + data += (-(10 ** (-(max_exponent_bits - 1) + i)) * means).tolist() + + data.append(0) + data.append(1.0) + + assert len(data) == 2**total_bits + + gap = 256 - len(data) + for i in range(gap): + data.append(0) + + data.sort() + return Tensor(data) + + +def weight_quantize( + A: torch.Tensor, + absmax: torch.Tensor, + out: torch.Tensor, + blocksize: int, + n: int, + quant_dtype: str, + code: torch.Tensor = None, +): + """ + Args: + A: (row ,col) torch.float16, torch.bfloat16, torch.float32 + absmax: (blocks) torch.float32 + blocks = n // blocksize, blocks += 1 if n % blocksize > 0 else 0 + quant_dtype: str + 目前可支持"int8"/"fp4"/"nf4" + blocksize: int + 目前只支持4096, 2048, 1024, 512, 256, 128, 64 + n: int + n = A.numel() + out: (row ,col) torch.int8 + code: torch.float32 + the quantization map + Returns: + out: (row ,col) torch.int8 + + """ + assert quant_dtype == "int8" or "fp4" or "nf4" + if code is None and quant_dtype == "int8": + code = create_dynamic_map().to(A.device) + if isinstance(A, torch.Tensor) and not A.requires_grad: + ops.infer.weight_quantize(A, absmax, out, blocksize, n, quant_dtype, code) + else: + raise NotImplementedError() diff --git a/ixformer_sdk/inference/functions/residual_bias.py b/ixformer_sdk/inference/functions/residual_bias.py new file mode 100644 index 00000000..9b337f0b --- /dev/null +++ b/ixformer_sdk/inference/functions/residual_bias.py @@ -0,0 +1,47 @@ +from typing import Union + +import ixformer._C as ops +import torch + +__all__ = ["residual_bias", "ref_residual_bias"] + + +def ref_residual_bias( + input: torch.Tensor, + residual: torch.Tensor, + bias: torch.Tensor = None, + alpha: float = 1, +): + if bias is not None: + output = residual.float() * alpha + input.float() + bias.float() + else: + output = residual.float() * alpha + input.float() + + return output.to(residual.dtype) + + +def residual_bias( + input: torch.Tensor, + residual: torch.Tensor, + bias: torch.Tensor = None, + alpha: float = 1, + output: torch.Tensor = None +): + """ + Args: + input: [batch_count, seq_len, hidden_size] or [batch_tokens, hidden_size] torch.half + residual: [batch_count, seq_len, hidden_size] or [batch_tokens, hidden_size] torch.half + bias: [hidden_size] torch.half + alpha: float + Returns: + output: [batch_count, seq_len, hidden_size] or [batch_tokens, hidden_size] torch.half + """ + if output is None: + output = torch.empty_like(input) + if alpha is None: + alpha = 1 + if bias is not None: + ops.train.add_residual_bias_forward(input, residual, bias, alpha, output) + else: + ops.train.add_residual_bias_forward(input, residual, alpha, output) + return output diff --git a/ixformer_sdk/inference/functions/rms_norm.py b/ixformer_sdk/inference/functions/rms_norm.py new file mode 100644 index 00000000..31f0f649 --- /dev/null +++ b/ixformer_sdk/inference/functions/rms_norm.py @@ -0,0 +1,134 @@ +from typing import Union + +import ixformer._C as ops +import torch +from torch.nn import init + +__all__ = ["ref_rms_norm", "rms_norm", "ref_residual_rms_norm", "residual_rms_norm"] + + +def ref_rms_norm( + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-5, + output: torch.Tensor = None, +): + dtype = input.dtype + input = input.float() + weight = weight.float() + rms_out = input * torch.rsqrt(input.pow(2).mean(-1, keepdim=True) + eps) + rms_out = rms_out * weight + rms_out = rms_out.to(dtype) + if output is not None: + output.copy_(rms_out) + else: + output = rms_out + return output + + +def rms_norm( + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-5, + output: torch.Tensor = None, +): + """ + This function is deprecated, please use residual_rms_norm. + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + eps: float32 + Returns: + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + """ + if output is None: + output = torch.empty_like(input) + + ops.infer.rms_norm(input, weight, output, None, eps) + + return output + + +def ref_residual_rms_norm( + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-5, + residual_alpha: float = 1.0, + residual: torch.Tensor = None, + residual_bias: torch.Tensor = None, + output: torch.Tensor = None, + residual_output: torch.Tensor = None, + is_post: bool = False, +): + dtype = input.dtype + + if residual_bias is not None: + input = input + residual_bias + + if residual is not None: + residual_output = torch.add( + input, residual * residual_alpha, out=residual_output + ) + input = input.float() + residual.float() * residual_alpha + else: + input = input.float() + weight = weight.float() + rms_out = input * torch.rsqrt(input.pow(2).mean(-1, keepdim=True) + eps) + rms_out = rms_out * weight + rms_out = rms_out.to(dtype) + if output is not None: + output.copy_(rms_out) + else: + output = rms_out + + if is_post and residual_output is not None: + residual_output = output + + return output, residual_output + + +def residual_rms_norm( + input: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-5, + residual_alpha: float = 1.0, + residual: torch.Tensor = None, + residual_bias: torch.Tensor = None, + output: torch.Tensor = None, + residual_output: torch.Tensor = None, + is_post: bool = False, +): + """ + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + eps: float32 + residual_alpha: float32 + residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + is_post: bool + Returns: + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on input. + residual_output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on residual. + """ + if residual is None: + if output is None: + output = torch.empty_like(input) + + ops.infer.rms_norm(input, weight, output, residual_bias, eps) + else: + ops.infer.residual_rms_norm( + input, + residual, + weight, + output, + residual_output, + residual_bias, + residual_alpha, + eps, + is_post, + ) + residual_output = residual_output if residual_output is not None else residual + output = output if output is not None else input + + return output, residual_output diff --git a/ixformer_sdk/inference/functions/scaled_dot_product_attention.py b/ixformer_sdk/inference/functions/scaled_dot_product_attention.py new file mode 100644 index 00000000..eda071fe --- /dev/null +++ b/ixformer_sdk/inference/functions/scaled_dot_product_attention.py @@ -0,0 +1,120 @@ +import math +from typing import List, Union + +import torch + +from .flash_attn import ixinfer_flash_attn_pad + +__all__ = ["scaled_dot_product_attention", "ref_scaled_dot_product_attention"] + + +def ref_scaled_dot_product_attention( + query: "torch.Tensor", + key: "torch.Tensor", + value: "torch.Tensor", + attn_mask=None, + dropout_p=0.0, + is_causal=False, +): + assert len(query.shape) >= 3 + assert len(key.shape) >= 3 + assert len(value.shape) >= 3 + batch_size = query.shape[0] + L = query.shape[-2] + S = key.shape[-2] + + if is_causal and attn_mask is not None: + raise RuntimeError() + + if attn_mask is None and is_causal is False: + attn_mask = torch.ones([batch_size, 1, 1, S]).bool().to(query.device) + + if is_causal: + attn_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0).to(query.device) + if attn_mask.dtype == torch.bool: + attn_mask = ( + torch.zeros_like(attn_mask).to(query.dtype).masked_fill(~attn_mask, -10000) + ) + # attn_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) if is_causal else attn_mask + # attn_mask = attn_mask.masked_fill(not attn_mask, -float('inf')) if attn_mask.dtype==torch.bool else attn_mask + attn_weight = torch.softmax( + (query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))) + attn_mask, dim=-1 + ) + attn_weight = torch.dropout(attn_weight, dropout_p, True) + return attn_weight.to(query.dtype) @ value + + +def scaled_dot_product_attention( + query: "torch.Tensor", + key: "torch.Tensor", + value: "torch.Tensor", + attn_mask=None, + dropout_p=0.0, + is_causal=False, +): + # 1. pytorch版本 query,key的head_dim相等， value可以不相等。但是我们实现的版本需要相等 + # 2. pytorch版本的 L, S可以不相等， 但是我们必须相等，并且为64的倍数 + # 3. pytorch支持 加上 mask + # 4. mask: 0 表示padding[与后端相反，所以注意如果用传进来的，取反转int；如果是自己生成，则直接生成后端的mask] + + """ + Args: + query: (N, ..., L, E) torch.float16, torch.bfloat16 + key: (N, ..., S, E) torch.float16, torch.bfloat16 + value: (N, ..., S, E) torch.float16, torch.bfloat16 + attn_mask: (N, ..., L, S) bool, torch.float32 + dropout_p: float32 + Dropout probability; if greater than 0.0, dropout is applied + is_causal: bool + If true, assumes causal attention masking and errors if both attn_mask and is_causal are set. + Returns: + Tensor: (N, ..., L, E) torch.float16, torch.bfloat16 + """ + + assert len(query.shape) >= 4, "len(query.shape) <4" + assert len(key.shape) >= 4, "len(key.shape) <4" + assert len(value.shape) >= 4, "len(value.shape) <4" + query_shape = list(query.shape) + key_shape = list(key.shape) + value_shape = list(value.shape) + + batch_size = query_shape[0] + q_seq_len = query_shape[-2] + kv_seq_len = key_shape[-2] + + if len(query_shape) > 4: + query.view(batch_size, -1, q_seq_len, query_shape[-1]) + if len(key_shape) > 4: + key.view(batch_size, -1, kv_seq_len, key_shape[-1]) + if len(value_shape) > 4: + value.view(batch_size, -1, kv_seq_len, value_shape[-1]) + + training = query.requires_grad + atten_scale = 1.0 / (query.size(-1) ** 0.5) + + if is_causal and attn_mask is not None: # 两者必须有有一个 + raise RuntimeError() + + head_num = query.size(1) + # 注意，training,inference都支持mask广播； + if training: + raise NotImplementedError("not support training!") + + else: # inference支持mask广播； + if attn_mask is None and is_causal is False: # mask 全0 + attn_mask = None + # attn_mask = ( + # torch.zeros([batch_size, 1, 1, kv_seq_len]).int().to(query.device) + # ) # 底层代码1代表mask，0代表保留 + elif is_causal: # mask 下三角是0，上三角是1 + attn_mask = ( + torch.ones([batch_size, 1, q_seq_len, kv_seq_len], dtype=torch.int) + .triu(diagonal=1) + .to(query.device) + ) # 上三角是1，下三角和对角线是0 + elif attn_mask is not None: # 外部传进来的，取反，转int + assert attn_mask.dtype == torch.bool or attn_mask.dtype == torch.float + assert attn_mask.dim() == 4 # 必须是4维 + if attn_mask.dtype == torch.bool: + attn_mask = (~attn_mask).int() # 取非操作，然后转int + return ixinfer_flash_attn_pad(query, key, value, attn_mask, None, atten_scale) diff --git a/ixformer_sdk/inference/functions/smoothquant.py b/ixformer_sdk/inference/functions/smoothquant.py new file mode 100644 index 00000000..17971100 --- /dev/null +++ b/ixformer_sdk/inference/functions/smoothquant.py @@ -0,0 +1,510 @@ +import ixformer._C as ops +import torch +import torch.nn.functional as NNF + +__all__ = [ + "ref_dynamic_scaled_quant_dynamic_int8", + "dynamic_scaled_quant_dynamic_int8", + "dynamic_scaled_quant_smoothquant", + "ref_silu_and_mul_smoothquant", + "silu_and_mul_smoothquant", + "ref_residual_rms_norm_dynamic_int8", + "residual_rms_norm_dynamic_int8", + "ref_residual_layer_norm_dynamic_int8", + "residual_layer_norm_dynamic_int8", + "ref_layer_norm_2sb_smoothquant", + "layer_norm_2sb_smoothquant", + "ref_residual_layer_norm_2sb_smoothquant", + "residual_layer_norm_2sb_smoothquant", +] + + +def ref_dynamic_scaled_quant_dynamic_int8( + input: torch.Tensor, + smooth_scales: torch.Tensor = None, + i8_output: torch.Tensor = None, + output_scales: torch.Tensor = None, +): + if i8_output is None: + i8_output = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales is None: + output_scales = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + + scales_shape = input.shape[:-1] + output = input.float() + if smooth_scales is not None: + output *= smooth_scales.view(1, -1) + amax_, _ = torch.max(torch.abs(output), dim=-1, keepdim=True) + scales = amax_ / 127.0 + output = output / scales + output = torch.clamp(torch.round(output), -127, 127).to(torch.int8) + + if i8_output is not None: + i8_output.copy_(output) + output = i8_output + if output_scales is not None: + output_scales.view(-1).copy_(scales.view(-1)) + scales = output_scales + + return output, scales.view(scales_shape) + + +def dynamic_scaled_quant_dynamic_int8( + input: torch.Tensor, + smooth_scales: torch.Tensor = None, + i8_output: torch.Tensor = None, + output_scales: torch.Tensor = None, +): + """ + Args: + input: (..., k) torch.float16,torch.bfloat16 + smooth_scales: (k) torch.float16,torch.bfloat16 + if smooth_scales is None, api is dynamic-per-token quantization. + Returns: + i8_output: (..., k) torch.int8 + output_scales: (...) torch.float32 + """ + if i8_output is None: + i8_output = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales is None: + output_scales = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + hidden_size = input.shape[-1] + + if smooth_scales is None: + ops.infer.scaled_int8_quant(i8_output, input, output_scales, 1) + return i8_output, output_scales + + ops.infer.dynamic_scaled_quant_smoothquant( + input.view(-1, hidden_size), + smooth_scales, + i8_output.view(-1, hidden_size), + output_scales, + ) + + return i8_output, output_scales + + +# For backward compatibility +dynamic_scaled_quant_smoothquant = dynamic_scaled_quant_dynamic_int8 + + +def ref_silu_and_mul_smoothquant( + input, smooth_scales, i8_output=None, output_scales=None +): + x1, x2 = input.chunk(chunks=2, dim=-1) + x = NNF.silu(x1) * x2 + + return ref_dynamic_scaled_quant_dynamic_int8( + x, smooth_scales, i8_output, output_scales + ) + + +def silu_and_mul_smoothquant(input, smooth_scales, i8_output=None, output_scales=None): + """ + Args: + input: (..., 2*k) torch.float16,torch.bfloat16 + smooth_scales: (k) torch.float16,torch.bfloat16 + if smooth_scales is None, api is dynamic-per-token quantization. + Returns: + i8_output: (..., k) torch.int8 + output_scales: (...) torch.float32 + """ + if i8_output is None: + output_shape = input.shape[:-1] + (input.shape[-1] // 2,) + i8_output = torch.empty(output_shape, dtype=torch.int8, device=input.device) + if output_scales is None: + output_scales = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + + ops.infer.silu_and_mul_smoothquant(i8_output, input, smooth_scales, output_scales) + + return i8_output, output_scales + + +def ref_residual_rms_norm_dynamic_int8( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor = None, + residual_bias: torch.Tensor = None, + eps: float = 1e-5, + smooth_scales: torch.Tensor = None, + output: torch.Tensor = None, + residual_output: torch.Tensor = None, + output_scales: torch.Tensor = None, + is_post: bool = False, +): + dtype = input.dtype + if output is None: + output = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales is None: + output_scales = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + if residual_bias is not None: + input = input + residual_bias + + if residual is not None: + residual_output = torch.add(input, residual, out=residual_output) + input = residual_output + input = input.float() + weight = weight.float() + rms_output = input * torch.rsqrt(input.pow(2).mean(-1, keepdim=True) + eps) + rms_output = (rms_output * weight).to(dtype) + + if residual is not None and is_post: + residual_output.copy_(rms_output) + output, output_scales = ref_dynamic_scaled_quant_dynamic_int8( + rms_output, smooth_scales, output, output_scales.view(-1) + ) + + return output, residual_output, output_scales + + +def residual_rms_norm_dynamic_int8( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor = None, + residual_bias: torch.Tensor = None, + eps: float = 1e-5, + smooth_scales: torch.Tensor = None, + output: torch.Tensor = None, + residual_output: torch.Tensor = None, + output_scales: torch.Tensor = None, + is_post: bool = False, +): + """ + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + eps: float32 + smooth_scales: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + is_post: bool + Returns: + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual_output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on residual. + output_scales: (...) torch.float16, torch.bfloat16, torch.float32 + """ + if output is None: + output = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales is None: + output_scales = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + + if residual is None: + ops.infer.rmsnorm_dynamic_int8( + input, weight, output, output_scales, smooth_scales, residual_bias, eps + ) + else: + ops.infer.residual_rmsnorm_dynamic_int8( + input, + residual, + weight, + output, + output_scales, + smooth_scales, + residual_output, + residual_bias, + eps, + is_post, + ) + residual_output = residual if residual_output is None else residual_output + + return output, residual_output, output_scales + + +def ref_residual_layer_norm_dynamic_int8( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + residual_bias: torch.Tensor = None, + eps: float = 1e-5, + smooth_scales: torch.Tensor = None, + output: torch.Tensor = None, + residual_output: torch.Tensor = None, + output_scales: torch.Tensor = None, +): + normalized_shape = [weight.size(-1)] + + if output is None: + output = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales is None: + output_scales = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + + if residual_bias is not None: + input = input + residual_bias + + if residual is not None: + residual_output = torch.add(input, residual, out=residual_output) + input = residual_output + + norm_output = torch.nn.functional.layer_norm( + input, normalized_shape, weight, bias, eps=eps + ) + + output, output_scales = ref_dynamic_scaled_quant_dynamic_int8( + norm_output, smooth_scales, output, output_scales.view(-1) + ) + + return output, residual_output, output_scales + + +def residual_layer_norm_dynamic_int8( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor = None, + residual_bias: torch.Tensor = None, + eps: float = 1e-5, + smooth_scales: torch.Tensor = None, + output: torch.Tensor = None, + residual_output: torch.Tensor = None, + output_scales: torch.Tensor = None, +): + """ + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + eps: float32 + smooth_scales: (hidden_size) torch.float16, torch.bfloat16, torch.float32 + Returns: + output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 + residual_output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on residual. + output_scales: (...) torch.float16, torch.bfloat16, torch.float32 + """ + if output is None: + output = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales is None: + output_scales = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + + if residual is None: + ops.infer.layer_norm_dynamic_int8( + input, + weight, + bias, + output, + output_scales, + smooth_scales, + residual_bias, + eps, + ) + else: + ops.infer.residual_layer_norm_dynamic_int8( + input, + residual, + weight, + bias, + output, + output_scales, + smooth_scales, + residual_output, + residual_bias, + eps, + ) + residual_output = residual_output if residual_output is not None else residual + + return output, residual_output, output_scales + + +def ref_layer_norm_2sb_smoothquant( + input, + weight1, + bias1, + smooth_scales1, + weight2, + bias2, + smooth_scales2, + i8_output1=None, + output_scales1=None, + i8_output2=None, + output_scales2=None, + eps=1e-5, +): + input1 = torch.nn.functional.layer_norm( + input, [weight1.shape[-1]], weight1, bias1, eps=eps + ) + input2 = torch.nn.functional.layer_norm( + input, [weight2.shape[-1]], weight2, bias2, eps=eps + ) + + i8_output1, output_scales1 = ref_dynamic_scaled_quant_dynamic_int8( + input1, smooth_scales1, i8_output1, output_scales1 + ) + i8_output2, output_scales2 = ref_dynamic_scaled_quant_dynamic_int8( + input2, smooth_scales2, i8_output2, output_scales2 + ) + + return i8_output1, output_scales1, i8_output2, output_scales2 + + +def layer_norm_2sb_smoothquant( + input, + weight1, + bias1, + smooth_scales1, + weight2, + bias2, + smooth_scales2, + output1=None, + output_scales1=None, + output2=None, + output_scales2=None, + eps=1e-5, +): + """ + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16 + weight1: (hidden_size) torch.float16, torch.bfloat16 + bias1: (hidden_size) torch.float16, torch.bfloat16 + smooth_scales1: (hidden_size) torch.float16, torch.bfloat16 + weight2: (hidden_size) torch.float16, torch.bfloat16 + bias2: (hidden_size) torch.float16, torch.bfloat16 + smooth_scales2: (hidden_size) torch.float16, torch.bfloat16 + eps: float32 + Returns: + output1: (..., hidden_size) torch.float16, torch.bfloat16 + output_scales1: (...) torch.float16, torch.bfloat16 + output2: (..., hidden_size) torch.float16, torch.bfloat16 + output_scales2: (...) torch.float16, torch.bfloat16 + """ + if output1 is None: + output1 = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales1 is None: + output_scales1 = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + if output2 is None: + output2 = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales2 is None: + output_scales2 = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + + ops.infer.layer_norm_2sb_smoothquant( + input, + weight1, + bias1, + smooth_scales1, + weight2, + bias2, + smooth_scales2, + output1, + output_scales1, + output2, + output_scales2, + eps, + ) + + return output1, output_scales1, output2, output_scales2 + + +def ref_residual_layer_norm_2sb_smoothquant( + input, + residual, + weight1, + bias1, + smooth_scales1, + weight2, + bias2, + smooth_scales2, + i8_output1=None, + output_scales1=None, + i8_output2=None, + output_scales2=None, + eps=1e-5, +): + residual_out = input + residual + input1 = torch.nn.functional.layer_norm( + residual_out, [weight1.shape[-1]], weight1, bias1, eps=eps + ) + input2 = torch.nn.functional.layer_norm( + residual_out, [weight2.shape[-1]], weight2, bias2, eps=eps + ) + + i8_output1, output_scales1 = ref_dynamic_scaled_quant_dynamic_int8( + input1, smooth_scales1, i8_output1, output_scales1 + ) + i8_output2, output_scales2 = ref_dynamic_scaled_quant_dynamic_int8( + input2, smooth_scales2, i8_output2, output_scales2 + ) + + return residual_out, i8_output1, output_scales1, i8_output2, output_scales2 + + +def residual_layer_norm_2sb_smoothquant( + input, + residual, + weight1, + bias1, + smooth_scales1, + weight2, + bias2, + smooth_scales2, + output1=None, + output_scales1=None, + output2=None, + output_scales2=None, + eps=1e-5, +): + """ + Args: + input: (..., hidden_size) torch.float16, torch.bfloat16 + residual: (..., hidden_size) torch.float16, torch.bfloat16 + weight1: (hidden_size) torch.float16, torch.bfloat16 + bias1: (hidden_size) torch.float16, torch.bfloat16 + smooth_scales1: (hidden_size) torch.float16, torch.bfloat16 + weight2: (hidden_size) torch.float16, torch.bfloat16 + bias2: (hidden_size) torch.float16, torch.bfloat16 + smooth_scales2: (hidden_size) torch.float16, torch.bfloat16 + eps: float32 + Returns: + output1: (..., hidden_size) torch.float16, torch.bfloat16 + output_scales1: (...) torch.float16, torch.bfloat16 + output2: (..., hidden_size) torch.float16, torch.bfloat16 + output_scales2: (...) torch.float16, torch.bfloat16 + """ + if output1 is None: + output1 = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales1 is None: + output_scales1 = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + if output2 is None: + output2 = torch.empty(input.shape, dtype=torch.int8, device=input.device) + if output_scales2 is None: + output_scales2 = torch.empty( + input.shape[:-1], dtype=torch.float32, device=input.device + ) + + ops.infer.residual_layer_norm_2sb_smoothquant( + input, + residual, + weight1, + bias1, + smooth_scales1, + weight2, + bias2, + smooth_scales2, + output1, + output_scales1, + output2, + output_scales2, + eps, + ) + + return residual, output1, output_scales1, output2, output_scales2 diff --git a/ixformer_sdk/inference/functions/softmax.py b/ixformer_sdk/inference/functions/softmax.py new file mode 100644 index 00000000..922443ef --- /dev/null +++ b/ixformer_sdk/inference/functions/softmax.py @@ -0,0 +1,33 @@ +from typing import Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = ["softmax", "ref_softmax"] + + +def ref_softmax(input: torch.Tensor, dim: int = None, _stacklevel: int = 3, dtype=None): + out = torch.nn.functional.softmax( + input, dim=dim, _stacklevel=_stacklevel, dtype=dtype + ) + return out + + +def softmax(input: torch.Tensor, dim=None, _stacklevel=3, dtype=None): + + """ + Args: + input: (...) torch.float16 + dim: int + 要进行softmax的维度,目前只支持最后一维， dim==-1 or dim == input.dim()-1 + _stacklevel: int + 这个参数只是为了与pytorch中对齐。 stacklevel is used in python to indicate warning mechanism how far up the stack it has to go to find the line that called the function which issued the warning. + dtype: torch.float16 + Returns: + Tensor: (...) torch.float16 + """ + output = torch.empty_like(input) + ops.infer.softmax(input, output, dim) + output = output.to(dtype) + return output diff --git a/ixformer_sdk/inference/functions/store_kv_cache.py b/ixformer_sdk/inference/functions/store_kv_cache.py new file mode 100644 index 00000000..112b90ae --- /dev/null +++ b/ixformer_sdk/inference/functions/store_kv_cache.py @@ -0,0 +1,75 @@ +import ixformer._C as ops +import torch + +__all__ = [ + "store_kv_cache", + "ref_store_kv_cache", +] + + +def ref_store_kv_cache( + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_batch_idx: torch.Tensor, + cache_seqlens: torch.Tensor, +): + """ + Args: + k: (batch_size, seqlen_new, head_num, head_dim) torch.float16, torch.bfloat16 + v: (batch_size, seqlen_new, head_num, head_dim) torch.float16, torch.bfloat16 + k_cache: (batch_size_cache, seqlen_cache, head_num, head_dim) torch.float16, torch.bfloat16 + v_cache: (batch_size_cache, seqlen_cache, head_num, head_dim) torch.float16, torch.bfloat16 + cache_batch_idx: (batch_size,) torch.int32 + The indices used to index into the KV cache. + cache_seqlens: (batch_size,) torch.int32 + The sequence lengths of the KV cache. + Returns: + None + """ + # 等价实现 + # concatenate k with k_cache, starting at the indices specified by cache_seqlens. + seqlen_new = k.size(1) + for kv_batch_idx, cache_kv_batch_idx in enumerate(cache_batch_idx): + cache_len = cache_seqlens[kv_batch_idx] + cache_start_idx = cache_len + cache_end_idx = cache_len + seqlen_new + + k_cache[cache_kv_batch_idx, cache_start_idx:cache_end_idx] = k[kv_batch_idx] + v_cache[cache_kv_batch_idx, cache_start_idx:cache_end_idx] = v[kv_batch_idx] + + +def store_kv_cache( + k: torch.Tensor, + v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_batch_idx: torch.Tensor, + cache_seqlens: torch.Tensor, +): + """ + Currently, only head_dim%2==0 is supported. + Args: + k: (batch_size, seqlen_new, head_num, head_dim) torch.float16, torch.bfloat16 + v: (batch_size, seqlen_new, head_num, head_dim) torch.float16, torch.bfloat16 + k_cache: (batch_size_cache, seqlen_cache, head_num, head_dim) torch.float16, torch.bfloat16 + v_cache: (batch_size_cache, seqlen_cache, head_num, head_dim) torch.float16, torch.bfloat16 + cache_batch_idx: (batch_size,) torch.int32 + The indices used to index into the KV cache. + cache_seqlens: (batch_size,) torch.int32 + The sequence lengths of the KV cache + Returns: + None + """ + head_dim = k.size(-1) + assert head_dim % 2 == 0, "Currently, only head_dim%2==0 is supported." + + ops.infer.store_kv_cache( + k, + v, + k_cache, + v_cache, + cache_batch_idx, + cache_seqlens, + ) diff --git a/ixformer_sdk/inference/functions/t5.py b/ixformer_sdk/inference/functions/t5.py new file mode 100644 index 00000000..ecfc2bc2 --- /dev/null +++ b/ixformer_sdk/inference/functions/t5.py @@ -0,0 +1,97 @@ +from typing import Union + +import ixformer._C as ops +import torch + +__all__ = [ + "t5_split_qkv", + "t5_split_qkv_update_kv_cache", + "ref_t5_split_qkv_update_kv_cache", + "ref_t5_split_qkv", +] + + +def reshape_query(query, head_num, head_dim): + batch_size, seq_len, _ = query.shape + query = query.view(batch_size, seq_len, head_num, head_dim) + query = query.transpose(1, 2) + return query + + +def ref_t5_split_qkv(qkv: "torch.Tensor", head_num: int, head_dim: int): + assert qkv.size(-1) == head_dim * head_num * 3 + batch_size, seq_len, _ = qkv.shape + q, k, v = torch.chunk(qkv, 3, dim=-1) + q = reshape_query(q, head_num, head_dim) + k = reshape_query(k, head_num, head_dim) + v = reshape_query(v, head_num, head_dim) + return q, k, v + + +def t5_split_qkv(qkv: "torch.Tensor", head_num: int, head_dim: int): + + """ + Args: + qkv: (batch_size, seq_len, head_dim * head_num * 3) torch.half, torch.bfloat16 + head_num: int + head_dim: int + Returns: + q: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16 + k: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16 + v: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16 + """ + batch_size, seq_len, _ = qkv.shape + q = qkv.new_empty([batch_size, head_num, seq_len, head_dim]) + k = qkv.new_empty([batch_size, head_num, seq_len, head_dim]) + v = qkv.new_empty([batch_size, head_num, seq_len, head_dim]) + ops.infer.t5_split_qkv(qkv, q, k, v, head_num, head_dim) + return q, k, v + + +def ref_t5_split_qkv_update_kv_cache( + qkv: "torch.Tensor", + past_key: "torch.Tensor", + past_value: "torch.Tensor", + head_num: int, + head_dim: int, +): + assert qkv.size(-1) == head_dim * head_num * 3 + batch_size, seq_len, _ = qkv.shape + q, k, v = torch.chunk(qkv, 3, dim=-1) + q = reshape_query(q, head_num, head_dim) + k = reshape_query(k, head_num, head_dim) + v = reshape_query(v, head_num, head_dim) + k = torch.cat([past_key, k], dim=2) + v = torch.cat([past_value, v], dim=2) + return q, k, v + + +def t5_split_qkv_update_kv_cache( + qkv: "torch.Tensor", + past_key: "torch.Tensor", + past_value: "torch.Tensor", + head_num: int, + head_dim: int, +): + + """ + Args: + qkv: (batch_size, 1 , head_dim * head_num * 3) torch.half, torch.bfloat16 + past_key: (batch_size, head_num, seq_len - 1, head_dim) torch.half, torch.bfloat16 + past_value: (batch_size, head_num, seq_len - 1, head_dim) torch.half, torch.bfloat16 + head_num: int + head_dim: int + Returns: + q: (batch_size, head_num, 1, head_dim) torch.half, torch.bfloat16 + k: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16 + v: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16 + """ + batch_size, _, past_seq_len, _ = list(past_key.shape) + seq_len = past_seq_len + 1 + q = qkv.new_empty([batch_size, head_num, 1, head_dim]) + k = qkv.new_empty([batch_size, head_num, seq_len, head_dim]) + v = qkv.new_empty([batch_size, head_num, seq_len, head_dim]) + ops.infer.t5_split_qkv_update_kv_cache( + qkv, past_key, past_value, q, k, v, head_num, head_dim + ) + return q, k, v diff --git a/ixformer_sdk/inference/functions/tgi.py b/ixformer_sdk/inference/functions/tgi.py new file mode 100644 index 00000000..3737e22a --- /dev/null +++ b/ixformer_sdk/inference/functions/tgi.py @@ -0,0 +1,617 @@ +import math +from typing import List, Optional + +import ixformer._C as ops +import torch + +__all__ = [ + "tgi_apply_rotary_emb_torch", + "tgi_apply_rotary", + "tgi_gather_prefill_logprobs", + "ref_paged_attention_v1", + "ref_paged_attention_v3", + "get_alibi_slopes", + "paged_attention_v1", + "reshape_and_cache_v1", + "paged_attention_v7", + "reshape_and_cache", + "paged_attention_v3", + "ref_reshape_and_cache_v3", + "reshape_and_cache_v3", +] + + +def get_alibi_slopes(total_num_heads: int) -> torch.Tensor: + closest_power_of_2 = 2 ** math.floor(math.log2(total_num_heads)) + base = torch.tensor( + 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), + dtype=torch.float32, + ) + powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32) + slopes = torch.pow(base, powers) + + if closest_power_of_2 != total_num_heads: + extra_base = torch.tensor( + 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), + dtype=torch.float32, + ) + num_remaining_heads = min( + closest_power_of_2, total_num_heads - closest_power_of_2 + ) + extra_powers = torch.arange( + start=1, end=1 + 2 * num_remaining_heads, step=2, dtype=torch.int32 + ) + slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0) + return slopes + + +def get_alibi_mask(num_heads, seqlen, device, dtype): + x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1) + y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1) + offsets = -(y - x).view(1, 1, seqlen) + return offsets + + +def ref_masked_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + attn_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + query = query * scale + dtype = query.dtype + device = query.device + query = query.to(torch.float32).cpu() + key = key.to(torch.float32).cpu() + value = value.to(torch.float32).cpu() + attn = torch.einsum("qhd,khd->hqk", query, key) + if attn_mask is not None: + attn_mask = attn_mask.cpu() + attn = attn + attn_mask + attn = torch.softmax(attn, dim=-1) + out = torch.einsum("hqk,khd->qhd", attn, value) + out = out.to(device).to(dtype) + return out + + +def ref_paged_attention_v1( + output: torch.Tensor, + query: torch.Tensor, + num_q_per_kv: int, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + use_alibi: bool, +) -> None: + num_query_heads = query.shape[1] + num_kv_heads = value_cache.shape[1] + head_size = value_cache.shape[2] + block_size = value_cache.shape[3] + + num_input_tokens = query.shape[0] + device = output.device + slopes = ( + get_alibi_slopes(num_query_heads) + .to(device) + .to(torch.float32) + .view(num_query_heads, 1, 1) + ) + for i in range(num_input_tokens): + q = query[i].unsqueeze(0) + block_table = block_tables[i] + context_len = int(context_lens[i]) + + keys = [] + values = [] + for j in range(context_len): + block_number = int(block_table[j // block_size]) + block_offset = j % block_size + + k = key_cache[block_number, :, :, block_offset, :] + k = k.reshape(num_kv_heads, head_size) + keys.append(k) + + v = value_cache[block_number, :, :, block_offset] + values.append(v) + keys = torch.stack(keys, dim=0) + values = torch.stack(values, dim=0) + if num_q_per_kv > 1: + keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1) + values = torch.repeat_interleave(values, num_q_per_kv, dim=1) + scale = 1.0 / (head_size**0.5) + if use_alibi: + offsets = get_alibi_mask( + num_query_heads, context_len, output.device, output.dtype + ) + mask = offsets * slopes + mask = mask.to(output.dtype) + else: + mask = None + out = ref_masked_attention( + q, + keys, + values, + scale, + mask, + ) + out = out.view(num_query_heads, head_size) + output[i].copy_(out, non_blocking=True) + + +def ref_paged_attention_v3( + output: torch.Tensor, + query: torch.Tensor, + num_q_per_kv: int, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + use_alibi: bool, +) -> None: + num_query_heads = query.shape[1] + num_kv_heads = value_cache.shape[1] + head_size = query.shape[2] + block_size = value_cache.shape[2] * 4 + + num_input_tokens = query.shape[0] + device = output.device + slopes = ( + get_alibi_slopes(num_query_heads) + .to(device) + .to(torch.float32) + .view(num_query_heads, 1, 1) + ) + for i in range(num_input_tokens): + q = query[i].unsqueeze(0) + block_table = block_tables[i] + context_len = int(context_lens[i]) + + keys = [] + values = [] + for j in range(context_len): + block_number = int(block_table[j // block_size]) + block_offset = j % block_size + + k = key_cache[block_number, :, block_offset // 4, :, block_offset % 4, :] + k = k.reshape(num_kv_heads, head_size) + keys.append(k) + + v = value_cache[block_number, :, block_offset // 4, :, block_offset % 4, :] + v = v.reshape(num_kv_heads, head_size) + values.append(v) + keys = torch.stack(keys, dim=0) + values = torch.stack(values, dim=0) + if num_q_per_kv > 1: + keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1) + values = torch.repeat_interleave(values, num_q_per_kv, dim=1) + scale = 1.0 / (head_size**0.5) + if use_alibi: + offsets = get_alibi_mask( + num_query_heads, context_len, output.device, output.dtype + ) + mask = offsets * slopes + mask = mask.to(output.dtype) + else: + mask = None + out = ref_masked_attention( + q, + keys, + values, + scale, + mask, + ) + out = out.view(num_query_heads, head_size) + output[i].copy_(out, non_blocking=True) + + +def rotate_half(x, interleaved=False): + if not interleaved: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + else: + x1, x2 = x[..., ::2], x[..., 1::2] + seq_len, head_nums, _ = x.shape + return torch.stack((-x2, x1), dim=-1).reshape(seq_len, head_nums, -1) + + +def tgi_apply_rotary_emb_torch( + x: "torch.Tensor", + cos: "torch.Tensor", + sin: "torch.Tensor", + interleaved: bool = False, +): + """ + x: (seqlen, num_heads, headdim) + cos, sin: (seqlen, 1, rotary_dim / 2) + interleaved: bool. 在interleaved的实现中,对奇偶维度旋转需要将维度两两交错，实现较为复杂。 + """ + ro_dim = cos.shape[-1] * 2 + assert ro_dim <= x.shape[-1] + assert cos.shape == sin.shape + if cos.dim() == 2: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + if interleaved: + cos = cos.repeat_interleave(2, dim=-1) + sin = sin.repeat_interleave(2, dim=-1) + else: + cos = cos.repeat(1, 1, 2) + sin = sin.repeat(1, 1, 2) + return torch.cat( + [ + x[..., :ro_dim].float() * cos.float() + + rotate_half(x[..., :ro_dim].float(), interleaved) * sin.float(), + x[..., ro_dim:].float(), + ], + dim=-1, + ).to(x.dtype) + + +def tgi_apply_rotary( + querys: List[torch.Tensor], + cos: "torch.Tensor", + sin: "torch.Tensor", + outs: List[torch.Tensor] = None, + is_neox_style: bool = True, +): + """ + Args: + querys: [(num_tokens, num_heads, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16] + cos: (max_position, 1, head_size //2) torch.half, torch.float, torch.bfloat16 + sin: (max_position, 1, head_size //2) torch.half, torch.float, torch.bfloat16 + is_neox_style: bool + 判断是否使用Neox,默认为True,即不使用interleaved + outs: [(num_tokens, num_heads, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16] + Returns: + outs: [(num_tokens, num_heads, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16] + """ + + assert sin.shape == cos.shape + rotary_dim = cos.shape[-1] + return_type = False + if len(querys) == 1: + query = querys[0] + query_dim = query.shape[-1] + query1 = query[..., :rotary_dim] + query2 = query[..., rotary_dim : 2 * rotary_dim] + elif len(querys) == 2: + return_type = True + query1 = querys[0] + query2 = querys[1] + assert query1.shape == query2.shape + query_dim = query1.shape[-1] * 2 + else: + raise ValueError( + f"Invalid number for querys: {len(querys)}. " "Expected number 1, or 2." + ) + + assert rotary_dim * 2 <= query_dim + if outs is None: + query_shape = query1.shape + out = torch.empty(*(query_shape[:-1] + [rotary_dim * 2])) + out1 = out[..., :rotary_dim] + out2 = out[..., rotary_dim : 2 * rotary_dim] + else: + assert len(querys) == len(outs) + for query, out in zip(querys, outs): + assert query.shape == out.shape + if len(outs) == 1: + out = outs[0] + out1 = out[..., :rotary_dim] + out2 = out[..., rotary_dim : 2 * rotary_dim] + else: + out1 = outs[0] + out2 = outs[1] + + if cos.dim() == 2: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + + ops.infer.tgi_rotary_embedding_neox( + query1, query2, cos, sin, out1, out2, is_neox_style + ) + + if return_type: + return out1, out2 + else: + return torch.cat([out1, out2], dim=-1) + + +def tgi_gather_prefill_logprobs( + logits: "torch.Tensor", + prefill_tokens_indices: "torch.Tensor", + output: "torch.Tensor" = None, +): + """ + Args: + logits: (num_tokens, vocab_size) torch.half, torch.bfloat16 + prefill_tokens_indices: (tokens_indices) torch.int + output: (tokens_indices, 1) torch.half, torch.bfloat16 + Returns: + output: (tokens_indices, 1) torch.half, torch.bfloat16 + """ + if output is None: + output = logits.new_empty(prefill_tokens_indices.shape) + ops.infer.tgi_gather_prefill_logprobs(logits, prefill_tokens_indices, output) + return output + + +def paged_attention_v1( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + use_sqrt_alibi: bool = False, +): + """ + Args: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + num_kv_heads: int + scale: float + block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64 + context_lens_cpu: (num_tokens) torch.int32 + context_lens: (num_tokens) torch.int32 + block_size: int + max_context_len: int + alibi_slopes: (num_heads) torch.float32 + use_sqrt_alibi: bool + Returns: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + """ + ops.infer.tgi_single_query_cached_kv_attention( + output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + context_lens, + block_size, + max_context_len, + query.stride(0), + use_sqrt_alibi, + alibi_slopes, + ) + + +def paged_attention_v3( + output: "torch.Tensor", + query: "torch.Tensor", + key_cache: "torch.Tensor", + value_cache: "torch.Tensor", + head_mapping: "torch.Tensor", + scale: float, + block_tables: "torch.Tensor", + context_lens: "torch.Tensor", + block_size: int, + max_context_len: int, + alibi_slopes: "torch.Tensor" = None, + use_sqrt_alibi: bool = False, +): + """ + Args: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + num_kv_heads: int + scale: float + block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64 + context_lens_cpu: (num_tokens) torch.int32 + context_lens: (num_tokens) torch.int32 + block_size: int + max_context_len: int + alibi_slopes: (num_heads) torch.float32 + use_sqrt_alibi: bool + Returns: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + """ + ops.infer.single_query_cached_kv_attention_v3( + output, + query, + key_cache, + value_cache, + head_mapping, + scale, + block_tables, + context_lens, + block_size, + max_context_len, + query.stride(0), + use_sqrt_alibi, + alibi_slopes, + ) + + + +def paged_attention_v7( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: int, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + use_sqrt_alibi: bool = False, +): + + """ + Args: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + num_kv_heads: int + scale: float + block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64 + context_lens: (num_tokens) torch.int32 + block_size: int + max_context_len: int + alibi_slopes: (num_heads) torch.float32 + use_sqrt_alibi: bool + Returns: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + """ + num_blocks = key_cache.size(0) + head_size = query.size(-1) + key_cache = key_cache.view(num_blocks, num_kv_heads, block_size, head_size) + value_cache = value_cache.view(num_blocks, num_kv_heads, block_size, head_size) + ops.infer.vllm_paged_attention( + output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + context_lens, + block_size, + max_context_len, + alibi_slopes, + True, + -1, + -1, + 0.0, + False, + use_sqrt_alibi, + ) + return output + +def reshape_and_cache_v1( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, +): + """ + Args: + key: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + value: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + key_cache: (num_blocks, num_heads, head_size//8, block_size, 8) or (num_blocks, num_heads, head_size//4, block_size, 4) torch.half, torch.float, torch.bfloat16 + if dtype=torch.half or torch.bfloat16,key_cache shape: (num_blocks, num_heads, head_size//8, block_size, 8) + if dtype=torch.float,key_cache shape: (num_blocks, num_heads, head_size//4, block_size, 4) + value_cache: (num_blocks, num_heads, head_size//8, block_size, 8) or (num_blocks, num_heads, head_size//4, block_size, 4) torch.half, torch.float, torch.bfloat16 + if dtype=torch.half or torch.bfloat16,value_cache shape: (num_blocks, num_heads, head_size//8, block_size, 8) + if dtype=torch.float,value_cache shape: (num_blocks, num_heads, head_size//4, block_size, 4) + slot_mapping: (num_tokens) torch.long + Returns: + None, 对key_cache,value_cache进行in place 操作 + """ + ops.infer.vllm_cache_ops_reshape_and_cache_v4( + key, + value, + key_cache, + value_cache, + slot_mapping, + key.stride(0), + value.stride(0), + ) + + +def reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, +): + """ + Args: + key: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + value: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + key_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.bfloat16 + value_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.bfloat16 + slot_mapping: (num_tokens) torch.long + Returns: + None, 对key_cache,value_cache进行in place 操作 + """ + num_tokens, num_kv_heads, head_size = key.shape + num_blocks = key_cache.size(0) + key_cache = key_cache.view(num_blocks, num_kv_heads, -1, head_size) + value_cache = value_cache.view(num_blocks, num_kv_heads, -1, head_size) + ops.infer.vllm_cache_ops_reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + key.stride(0), + value.stride(0), + ) + + +def ref_reshape_and_cache_v3( + key, + value, + key_cache, + value_cache, + slot_mapping, + num_tokens, + num_heads, + head_size, + block_size, +): + reshaped_key = key.view(num_tokens, num_heads, head_size // 32, 32) + reshaped_value = value.reshape(num_tokens, num_heads, head_size // 32, 32) + for i in range(num_tokens): + + block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor") + block_offset = slot_mapping[i] % block_size + + key_cache[ + block_idx, :, block_offset // 4, :, block_offset % 4, : + ] = reshaped_key[i] + value_cache[ + block_idx, :, block_offset // 4, :, block_offset % 4, : + ] = reshaped_value[i] + + +def reshape_and_cache_v3( + key: "torch.Tensor", + value: "torch.Tensor", + key_cache: "torch.Tensor", + value_cache: "torch.Tensor", + slot_mapping: "torch.Tensor", +): + """ + Args: + key: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + value: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + key_cache: (num_blocks, num_heads, block_size // 4, head_size // 32, 4, 32) torch.half, torch.bfloat16 + 目前block_size 只支持16,head_size 只支持64,128,256 + value_cache: (num_blocks, num_heads, block_size // 4, head_size // 32, 4, 32) torch.half, torch.bfloat16 + slot_mapping: (num_tokens) torch.int + Returns: + None, 对key_cache,value_cache进行in place 操作 + """ + if key.dim() != 3 or key.shape != value.shape or key.size(-1) not in [64, 128, 256]: + raise NotImplementedError( + "reshape_and_cache_v3 only support key.dim()==3 and key.shape== value.shape and head_size must be 64, 128 , 256!" + ) + ops.infer.cache_ops_reshape_and_cache_v3( + key, + value, + key_cache, + value_cache, + slot_mapping, + key.stride(0), + value.stride(0), + ) diff --git a/ixformer_sdk/inference/functions/vllm.py b/ixformer_sdk/inference/functions/vllm.py new file mode 100644 index 00000000..4edf1dd4 --- /dev/null +++ b/ixformer_sdk/inference/functions/vllm.py @@ -0,0 +1,2033 @@ +import math +from typing import Optional, Union + +import ixformer._C as ops +import ixformer._C._functions as CF +import torch + +from ixformer.core import config + +from .linear import linear +from .paged_attention import paged_attention as paged_attention_ixformer_impl + +__all__ = [ + "ref_vllm_paged_attention", + "vllm_paged_attention", + "ref_vllm_paged_attention_mla", + "vllm_paged_attention_mla", + "vllm_paged_attention_mla_fused", + "ref_vllm_paged_attention_mla_int8", + "vllm_paged_attention_mla_int8", + "ref_vllm_paged_attention_v5", + "vllm_paged_attention_v5", + "ref_vllm_paged_attention_v4", + "vllm_paged_attention_v4", + "ref_vllm_reshape_and_cache_v4", + "vllm_reshape_and_cache_v4", + "ref_vllm_reshape_and_cache", + "vllm_reshape_and_cache", + "vllm_cache_ops_reshape_and_cache", + "ref_reshape_and_cache_flash", + "reshape_and_cache_flash", + "ref_vllm_rotary_embedding", + "vllm_rotary_embedding", + "ref_vllm_rotary_embedding_phi", + "vllm_rotary_embedding_phi", + "ref_vllm_batched_rotary_embedding", + "vllm_batched_rotary_embedding", + "ref_vllm_copy_blocks", + "vllm_copy_blocks", + "ref_vllm_swap_blocks", + "vllm_swap_blocks", + "vllm_gather_cache", + "vllm_gather_cache_int8", + "ref_vllm_gather_cache_int8", + "ref_vllm_gather_cache", + "ref_vllm_concat_and_cache_mla", + "vllm_concat_and_cache_mla", + "ref_vllm_concat_and_cache_mla_int8", + "vllm_concat_and_cache_mla_int8", + "vllm_llama_mlp", + "gptq_gemm", + "vllm_gptq_shuffle", + "vllm_moe_topk_softmax", + "vllm_moe_align_block_size", + "ref_vllm_invoke_fused_moe_kernel", + "vllm_invoke_fused_moe_kernel", + "advance_step_flashattn", + "weak_ref_tensor", + # customized ops + "vllm_rotary_embedding_with_key_layer_norm", + "ref_vllm_rotary_embedding_with_key_layer_norm", +] + +weak_ref_tensor = ops.infer.weak_ref_tensor + + +def ref_vllm_paged_attention( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + softcap: float = 0.0, + window_left: int = -1, + window_right: int = -1, + use_sqrt_alibi: bool = False, +): + assert window_right in [-1, 0] + + def get_alibi_mask(num_heads, seqlen, device, dtype): + x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1) + y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1) + offsets = -(y - x).view(1, 1, seqlen) + return offsets + + def ref_masked_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + query = query * scale + dtype = query.dtype + device = query.device + query = query.to(torch.float32) + key = key.to(torch.float32) + value = value.to(torch.float32) + attn = torch.einsum("qhd,khd->hqk", query, key) + if attn_mask is not None: + attn_mask = attn_mask + attn = attn + attn_mask + attn = torch.softmax(attn, dim=-1) + out = torch.einsum("hqk,khd->qhd", attn, value) + out = out.to(device).to(dtype) + return out + + head_size = query.shape[-1] + num_query_heads = query.shape[1] + num_kv_heads = value_cache.shape[1] + num_input_tokens = query.shape[0] + + num_q_per_kv = num_query_heads // num_kv_heads + slopes = ( + alibi_slopes.view(num_query_heads, 1, 1) + if alibi_slopes is not None + else alibi_slopes + ) + + for i in range(num_input_tokens): + q = query[i].unsqueeze(0) + block_table = block_tables[i] + context_len = int(context_lens[i]) + + keys = [] + values = [] + for j in range(context_len): + block_number = int(block_table[j // block_size]) + block_offset = j % block_size + + k = key_cache[block_number, :, block_offset, :] + keys.append(k) + + v = value_cache[block_number, :, block_offset, :] + values.append(v) + keys = torch.stack(keys, dim=0) + values = torch.stack(values, dim=0) + if num_q_per_kv > 1: + keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1) + values = torch.repeat_interleave(values, num_q_per_kv, dim=1) + if alibi_slopes is not None: + offsets = get_alibi_mask( + num_query_heads, context_len, output.device, output.dtype + ) + mask = offsets * slopes + mask = mask.to(output.dtype) + if window_left != -1: + index = torch.ones_like(mask, dtype=torch.int32, device=mask.device) + index[:, :, (context_len - 1 - window_left) :] = 0 + index = index.bool() + mask.masked_fill_(index, float("-inf")) + else: + if window_left != -1: + mask = torch.zeros([1, 1, context_len], dtype=q.dtype, device=q.device) + index = torch.ones_like(mask, dtype=torch.int32, device=mask.device) + index[:, :, (context_len - 1 - window_left) :] = 0 + index = index.bool() + mask.masked_fill_(index, float("-inf")) + else: + mask = None + + out = ref_masked_attention( + q, + keys, + values, + scale, + mask, + ) + out = out.view(num_query_heads, head_size) + if softcap != 0.0: + out = softcap * torch.tanh(out / softcap) + output[i].copy_(out, non_blocking=True) + + return output + + +def vllm_paged_attention_ixinfer( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + softcap: float = 0.0, + causal: bool = True, + window_left: int = -1, + window_right: int = -1, + use_cuda_graph: bool = False, + use_sqrt_alibi: bool = False, +): + """ + Arguments: + query: [torch.half, torch.bfloat16] [num_tokens, num_heads, head_size] + key_cache: [torch.half, torch.bfloat16] [num_blocks, num_kv_heads, block_size, head_size] + value_cache: [torch.half, torch.bfloat16] [num_blocks, num_kv_heads, block_size, head_size] + num_kv_heads: int + scale: float + block_tables: [torch.int64] [num_tokens, max_num_blocks_per_seq] + context_lens: [torch.int32] [num_tokens] + block_size: int + max_context_len: int + alibi_slopes: [torch.float32] [num_heads] + softcap: float + causal: bool + window_left: int + window_right: int + use_sqrt_alibi: bool: False + Return: + output: [torch.half, torch.bfloat16] [num_tokens, num_heads, head_size] + """ + ops.infer.vllm_paged_attention( + output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + context_lens, + block_size, + max_context_len, + alibi_slopes, + causal, + window_left, + window_right, + softcap, + use_cuda_graph, + use_sqrt_alibi, + ) + return output + + +def vllm_paged_attention_ixformer( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + softcap: float = 0.0, + use_sqrt_alibi: bool = False, + need_view: bool = True, +): + """ + Args: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + num_kv_heads: int + scale: float + block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64 + context_lens_cpu: (num_tokens) torch.int32 + context_lens: (num_tokens) torch.int32 + block_size: int + max_context_len: int + alibi_slopes: (num_heads) torch.float32 + use_sqrt_alibi: bool + Returns: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + """ + + if need_view: + num_blocks = key_cache.size(0) + head_size = query.size(-1) + key_cache = key_cache.view(num_blocks, num_kv_heads, block_size, head_size) + value_cache = value_cache.view(num_blocks, num_kv_heads, block_size, head_size) + paged_attention_ixformer_impl( + output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + context_lens, + block_size, + max_context_len, + alibi_slopes, + use_sqrt_alibi, + ) + return output + + +def ref_vllm_paged_attention_mla( + output: torch.Tensor, + query: torch.Tensor, + kv_cache: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + max_context_len: int, +): + def ref_masked_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + query = query * scale + dtype = query.dtype + device = query.device + attn = torch.einsum("qhd,khd->hqk", query, key) + if attn_mask is not None: + attn_mask = attn_mask + attn = attn + attn_mask + attn = attn.to(torch.float) + attn = torch.softmax(attn, dim=-1) + value = value.to(torch.float) + out = torch.einsum("hqk,khd->qhd", attn, value) + out = out.to(device).to(dtype) + return out + + num_heads = query.shape[-2] + kv_lora_rank = output.shape[-1] + block_size = kv_cache.shape[1] + num_input_tokens = query.shape[0] + + for i in range(num_input_tokens): + q = query[i].unsqueeze(0) + block_table = block_tables[i] + context_len = int(context_lens[i]) + + keys = [] + values = [] + for j in range(context_len): + block_number = int(block_table[j // block_size]) + block_offset = j % block_size + + k = kv_cache[block_number, block_offset, :] + keys.append(k) + + v = kv_cache[block_number, block_offset, :kv_lora_rank] + values.append(v) + keys = torch.stack(keys, dim=0).unsqueeze(-2).repeat(1, num_heads, 1) + values = torch.stack(values, dim=0).unsqueeze(-2).repeat(1, num_heads, 1) + mask = None + + out = ref_masked_attention( + q, + keys, + values, + scale, + mask, + ) + out = out.view(num_heads, kv_lora_rank) + output[i].copy_(out, non_blocking=True) + + return output + + +def ref_vllm_paged_attention_mla_int8( + output: torch.Tensor, + query: torch.Tensor, + query_scale: torch.Tensor, + kv_cache: torch.Tensor, + kv_cache_scale: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + max_context_len: int, +): + def ref_masked_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + query = query * scale + dtype = query.dtype + device = query.device + attn = torch.einsum("qhd,khd->hqk", query, key) + if attn_mask is not None: + attn_mask = attn_mask + attn = attn + attn_mask + attn = attn.to(torch.float) + attn = torch.softmax(attn, dim=-1) + value = value.to(torch.float) + out = torch.einsum("hqk,khd->qhd", attn, value) + out = out.to(device).to(dtype) + return out + + # dequant q + num_heads = query.shape[-2] + kv_lora_rank = output.shape[-1] + block_size = kv_cache.shape[1] + num_input_tokens = query.shape[0] + query = query * query_scale.unsqueeze(-1) + + for i in range(num_input_tokens): + q = query[i].unsqueeze(0) + block_table = block_tables[i] + context_len = int(context_lens[i]) + + keys = [] + values = [] + for j in range(context_len): + block_number = int(block_table[j // block_size]) + block_offset = j % block_size + + k = kv_cache[block_number, block_offset, :kv_lora_rank] + k_scale = kv_cache_scale[block_number, block_offset, 0] + k_pe = kv_cache[block_number, block_offset, kv_lora_rank:] + k_pe_scale = kv_cache_scale[block_number, block_offset, 1] + k = k * k_scale + v = k + k_pe = k_pe * k_pe_scale + k = torch.cat((k, k_pe), dim=-1) + keys.append(k) + values.append(v) + keys = torch.stack(keys, dim=0).unsqueeze(-2).repeat(1, num_heads, 1) + values = torch.stack(values, dim=0).unsqueeze(-2).repeat(1, num_heads, 1) + mask = None + + out = ref_masked_attention( + q, + keys, + values, + scale, + mask, + ) + out = out.view(num_heads, kv_lora_rank) + output[i].copy_(out, non_blocking=True) + + return output + + +def vllm_paged_attention_mla( + output: torch.Tensor, + query: torch.Tensor, + kv_cache: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + max_context_len: int, + use_cuda_graph: bool = False, +): + """ + Args: + output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16 + query: (num_tokens, num_heads, kv_lora_rank+qk_rope_head_dim) torch.half, torch.bfloat16 + kv_cache: (num_blocks, block_size, kv_lora_rank+qk_rope_head_dim) torch.half, torch.bfloat16 + scale: float + block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64 + context_lens: (num_tokens) torch.int32 + max_context_len: int + use_cuda_graph: bool + Returns: + output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16 + """ + ops.infer.vllm_paged_attention_mla( + output, + query, + kv_cache, + scale, + block_tables, + context_lens, + max_context_len, + use_cuda_graph, + ) + return output + + +def vllm_paged_attention_mla_int8( + output: torch.Tensor, + query: torch.Tensor, + query_scale: torch.Tensor, + kv_cache: torch.Tensor, + kv_cache_scale: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + max_context_len: int, + use_cuda_graph: bool = False, +): + """ + Args: + output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16 + query: (num_tokens, num_heads, kv_lora_rank+qk_rope_head_dim) torch.int8 + query_scale: (num_tokens, num_heads) torch.float + kv_cache: (num_blocks, block_size, kv_lora_rank+qk_rope_head_dim) torch.half, torch.bfloat16 + kv_cache_scale: (num_blocks, block_size, 2) torch.float + scale: float + block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64 + context_lens: (num_tokens) torch.int32 + max_context_len: int + use_cuda_graph: bool + Returns: + output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16 + """ + ops.infer.vllm_paged_attention_mla_int8( + output, + query, + query_scale, + kv_cache, + kv_cache_scale, + scale, + block_tables, + context_lens, + max_context_len, + use_cuda_graph, + ) + return output + + +def vllm_paged_attention_mla_fused( + output: torch.Tensor, + q_nope: torch.Tensor, + q_pe: torch.Tensor, + kv_cache: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens: torch.Tensor, + max_context_len: int, + k_c_normed: torch.Tensor = None, + k_pe: torch.Tensor = None, + use_cuda_graph: bool = False, +): + """ + Args: + q_nope: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16 + q_pe: (num_tokens, num_heads, qk_rope_head_dim) torch.half, torch.bfloat16 + kv_cache: (num_blocks, block_size, kv_lora_rank+qk_rope_head_dim) torch.half, torch.bfloat16 + scale: float + block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64 + context_lens: (num_tokens) torch.int32 + max_context_len: int + k_c_normed: (num_tokens, kv_lora_rank) torch.half, torch.bfloat16 + k_pe: (num_tokens, qk_rope_head_dim) torch.half, torch.bfloat16 + use_cuda_graph: bool + Returns: + output: (num_tokens, num_heads, kv_lora_rank) torch.half, torch.bfloat16 + """ + ops.infer.vllm_paged_attention_mla_fused( + output, + q_nope, + q_pe, + kv_cache, + scale, + block_tables, + context_lens, + max_context_len, + k_c_normed, + k_pe, + use_cuda_graph, + ) + return output + + +def ref_vllm_paged_attention_v5( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens_cpu: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, +): + def get_alibi_mask(num_heads, seqlen, device, dtype): + x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1) + y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1) + offsets = -(y - x).view(1, 1, seqlen) + return offsets + + def ref_masked_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + query = query * scale + dtype = query.dtype + device = query.device + query = query.to(torch.float32).cpu() + key = key.to(torch.float32).cpu() + value = value.to(torch.float32).cpu() + attn = torch.einsum("qhd,khd->hqk", query, key) + if attn_mask is not None: + attn_mask = attn_mask.cpu() + attn = attn + attn_mask + attn = torch.softmax(attn, dim=-1) + out = torch.einsum("hqk,khd->qhd", attn, value) + out = out.to(device).to(dtype) + return out + + head_size = query.shape[-1] + num_query_heads = query.shape[1] + num_kv_heads = value_cache.shape[1] + num_input_tokens = query.shape[0] + + num_q_per_kv = num_query_heads // num_kv_heads + slopes = ( + alibi_slopes.view(num_query_heads, 1, 1) + if alibi_slopes is not None + else alibi_slopes + ) + + for i in range(num_input_tokens): + q = query[i].unsqueeze(0) + block_table = block_tables[i] + context_len = int(context_lens[i]) + + keys = [] + values = [] + for j in range(context_len): + block_number = int(block_table[j // block_size]) + block_offset = j % block_size + + k = key_cache[block_number, :, block_offset, :] + keys.append(k) + + v = value_cache[block_number, :, block_offset, :] + values.append(v) + keys = torch.stack(keys, dim=0) + values = torch.stack(values, dim=0) + if num_q_per_kv > 1: + keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1) + values = torch.repeat_interleave(values, num_q_per_kv, dim=1) + if alibi_slopes is not None: + offsets = get_alibi_mask( + num_query_heads, context_len, output.device, output.dtype + ) + mask = offsets * slopes + mask = mask.to(output.dtype) + else: + mask = None + out = ref_masked_attention( + q, + keys, + values, + scale, + mask, + ) + out = out.view(num_query_heads, head_size) + output[i].copy_(out, non_blocking=True) + + return output + + +def vllm_paged_attention_v5( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens_cpu: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + use_sqrt_alibi: bool = False, + need_view: bool = True, +): + """ + Args: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + num_kv_heads: int + scale: float + block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64 + context_lens_cpu: (num_tokens) torch.int32 + context_lens: (num_tokens) torch.int32 + block_size: int + max_context_len: int + alibi_slopes: (num_heads) torch.float32 + use_sqrt_alibi: bool + Returns: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + """ + if need_view: + num_blocks = key_cache.size(0) + head_size = query.size(-1) + key_cache = key_cache.view(num_blocks, num_kv_heads, block_size, head_size) + value_cache = value_cache.view(num_blocks, num_kv_heads, block_size, head_size) + ops.infer.vllm_paged_attention_v5( + output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + context_lens_cpu, + context_lens, + block_size, + max_context_len, + alibi_slopes, + use_sqrt_alibi, + ) + return output + + +def ref_vllm_paged_attention_v4( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens_cpu: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + use_sqrt_alibi: bool = False, +): + def get_alibi_mask(num_heads, seqlen, device, dtype): + x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1) + y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1) + offsets = -(y - x).view(1, 1, seqlen) + return offsets + + def ref_masked_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + scale: float, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + query = query * scale + dtype = query.dtype + device = query.device + query = query.to(torch.float32).cpu() + key = key.to(torch.float32).cpu() + value = value.to(torch.float32).cpu() + attn = torch.einsum("qhd,khd->hqk", query, key) + if attn_mask is not None: + attn_mask = attn_mask.cpu() + attn = attn + attn_mask + attn = torch.softmax(attn, dim=-1) + out = torch.einsum("hqk,khd->qhd", attn, value) + out = out.to(device).to(dtype) + return out + + num_query_heads = query.shape[1] + num_kv_heads = value_cache.shape[1] + head_size = value_cache.shape[2] + block_size = value_cache.shape[3] + num_input_tokens = query.shape[0] + + num_q_per_kv = num_query_heads // num_kv_heads + slopes = ( + alibi_slopes.view(num_query_heads, 1, 1) + if alibi_slopes is not None + else alibi_slopes + ) + + for i in range(num_input_tokens): + q = query[i].unsqueeze(0) + block_table = block_tables[i] + context_len = int(context_lens[i]) + + keys = [] + values = [] + for j in range(context_len): + block_number = int(block_table[j // block_size]) + block_offset = j % block_size + + k = key_cache[block_number, :, :, block_offset, :] + k = k.reshape(num_kv_heads, head_size) + keys.append(k) + + v = value_cache[block_number, :, :, block_offset] + values.append(v) + keys = torch.stack(keys, dim=0) + values = torch.stack(values, dim=0) + if num_q_per_kv > 1: + keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1) + values = torch.repeat_interleave(values, num_q_per_kv, dim=1) + scale = 1.0 / (head_size**0.5) + if alibi_slopes is not None: + offsets = get_alibi_mask( + num_query_heads, context_len, output.device, output.dtype + ) + mask = offsets * slopes + mask = mask.to(output.dtype) + else: + mask = None + out = ref_masked_attention( + q, + keys, + values, + scale, + mask, + ) + out = out.view(num_query_heads, head_size) + output[i].copy_(out, non_blocking=True) + + return output + + +def vllm_paged_attention_v4( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + num_kv_heads: torch.Tensor, + scale: float, + block_tables: torch.Tensor, + context_lens_cpu: torch.Tensor, + context_lens: torch.Tensor, + block_size: int, + max_context_len: int, + alibi_slopes: torch.Tensor = None, + use_sqrt_alibi: bool = False, +): + """ + Args: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16 + num_kv_heads: int + scale: float + block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64 + context_lens_cpu: (num_tokens) torch.int32 + context_lens: (num_tokens) torch.int32 + block_size: int + max_context_len: int + alibi_slopes: (num_heads) torch.float32 + use_sqrt_alibi: bool + Returns: + output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + """ + ops.infer.vllm_paged_attention_v4( + output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + context_lens_cpu, + context_lens, + block_size, + max_context_len, + alibi_slopes, + use_sqrt_alibi, + ) + return output + + +def ref_vllm_rotary_embedding( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + is_neox_style: bool = True, +): + def _rotate_neox(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _rotate_gptj(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., ::2] + x2 = x[..., 1::2] + x = torch.stack((-x2, x1), dim=-1) + return x.flatten(-2) + + query_shape = query.shape + key_shape = key.shape + B = query.shape[0] + query = query.view(B, -1, head_size) + key = key.view(B, -1, head_size) + + cos_sin = cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + if is_neox_style: + cos = cos.repeat(1, 1, 2).unsqueeze(-2) + sin = sin.repeat(1, 1, 2).unsqueeze(-2) + else: + cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) + sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) + + rotate_fn = _rotate_neox if is_neox_style else _rotate_gptj + query_rot = query * cos + rotate_fn(query) * sin + key_rot = key * cos + rotate_fn(key) * sin + + query = query_rot.flatten(-2).view(query_shape) + key = key_rot.flatten(-2).view(key_shape) + return query, key + + +def vllm_rotary_embedding( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + is_neox_style: bool = True, +): + + """ + Args: + positions: (num_tokens) torch.long + query: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16 + key: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16 + head_size: int + cos_sin_cache: (max_position, head_size) torch.half, torch.float, torch.bfloat16 + is_neox_style: bool + Returns: + None. 对query, key 做in place 操作 + """ + ops.infer.vllm_rotary_embedding( + positions, + query, + key, + head_size, + cos_sin_cache, + is_neox_style, + ) + + +def ref_vllm_rotary_embedding_phi( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + long_offset: torch.Tensor, + k: int, + offsets: torch.Tensor = None, +): + def _rotate_neox(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + query_shape = query.shape + key_shape = key.shape + B = query.shape[0] + query = query.view(B, -1, head_size) + key = key.view(B, -1, head_size) + + if long_offset is None: + long_offset = ( + torch.any(positions > k).float() * torch.full_like(positions, k) + ).long() + idx = torch.add(positions, long_offset) if long_offset is not None else positions + idx = torch.add(idx, offsets) if offsets is not None else idx + cos_sin = torch.index_select(cos_sin_cache, 0, idx) + + cos, sin = cos_sin.chunk(2, dim=-1) + cos = cos.repeat(1, 2).unsqueeze(-2) + sin = sin.repeat(1, 2).unsqueeze(-2) + + query = query * cos + _rotate_neox(query) * sin + key = key * cos + _rotate_neox(key) * sin + + query = query.flatten(-2).view(query_shape) + key = key.flatten(-2).view(key_shape) + + return query, key + + +def vllm_rotary_embedding_phi( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + long_offset: torch.Tensor, + k: int, + offsets: torch.Tensor = None, +): + """ + Args: + positions: (num_tokens) torch.long + query: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16 + key: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16 + cos_sin_cache: (max_position, head_size) torch.half, torch.float, torch.bfloat16 + head_size: int + cos_sin_cache: (max_position, head_size) torch.half, torch.float, torch.bfloat16 + long_offset: (1,) torch.bool + k: int + offsets: (num_tokens) torch.half, torch.float, torch.bfloat16 + Returns: + None. 对query, key 做in place 操作 + """ + ops.infer.vllm_rotary_embedding_phi( + positions, + query, + key, + head_size, + cos_sin_cache, + long_offset, + k, + offsets, + ) + + +def ref_vllm_rotary_embedding_with_key_layer_norm( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + key_out: torch.Tensor = None, + eps: float = 1e-5, + is_neox_style: bool = True, +): + B = key.shape[0] + query_size = query.size() + query, key = ref_vllm_rotary_embedding( + positions, + query.view(B, -1), + key.view(B, -1), + head_size, + cos_sin_cache, + is_neox_style, + ) + query = query.view(query_size) + key = key.view(B, -1, head_size) + + norm_key = torch.nn.functional.layer_norm( + key, + [ + head_size, + ], + weight, + bias, + eps, + ) + if key_out is not None: + key_out.copy_(norm_key) + else: + key_out = norm_key + return query, key_out + + +def vllm_rotary_embedding_with_key_layer_norm( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + key_out: torch.Tensor = None, + eps: float = 1e-5, + is_neox_style: bool = True, +): + """ + Args: + positions: (num_tokens) torch.int64 + query: (num_tokens, num_heads * head_size) or (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + key: (num_tokens, num_kv_heads * head_size) or (num_tokens, num_kv_heads, head_size) torch.half, torch.float, torch.bfloat16 + weight: (head_size) torch.half, torch.float, torch.bfloat16 + bias: (head_size) torch.half, torch.float, torch.bfloat16 + head_size: int + cos_sin_cache: (max_position, rot_dim) torch.half, torch.float, torch.bfloat16 + key_out: (num_tokens, num_kv_heads * head_size) or (num_tokens, num_kv_heads, head_size) torch.half, torch.float, torch.bfloat16 + eps: float + is_neox_style: bool + Returns: + query: (num_tokens, num_heads * head_size) or (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + key_out: (num_tokens, num_kv_heads * head_size) or (num_tokens, num_kv_heads, head_size) torch.half, torch.float, torch.bfloat16 + """ + ops.infer.vllm_rotary_embedding_with_key_layer_norm( + positions, + query, + key, + weight, + bias, + head_size, + cos_sin_cache, + key_out, + eps, + is_neox_style, + ) + key_out = key if key_out is None else key_out + return query, key_out + + +def ref_vllm_batched_rotary_embedding( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + is_neox_style, + rotary_dim: int, + offsets: torch.Tensor, +): + def _rotate_neox(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _rotate_gptj(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., ::2] + x2 = x[..., 1::2] + x = torch.stack((-x2, x1), dim=-1) + return x.flatten(-2) + + query = query.view(*query.shape[:-1], -1, head_size) + key = key.view(*key.shape[:-1], -1, head_size) + + query_rot = query[..., :rotary_dim] + key_rot = key[..., :rotary_dim] + if rotary_dim < head_size: + query_pass = query[..., rotary_dim:] + key_pass = key[..., rotary_dim:] + + cos_sin = cos_sin_cache[torch.add(positions, offsets)] + cos, sin = cos_sin.chunk(2, dim=-1) + if is_neox_style: + # NOTE(woosuk): Here we assume that the positions tensor has the + # shape [batch_size, seq_len]. + cos = cos.repeat(1, 1, 2).unsqueeze(-2) + sin = sin.repeat(1, 1, 2).unsqueeze(-2) + else: + cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2) + sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2) + + rotate_fn = _rotate_neox if is_neox_style else _rotate_gptj + query_rot = query_rot * cos + rotate_fn(query_rot) * sin + key_rot = key_rot * cos + rotate_fn(key_rot) * sin + + if rotary_dim < head_size: + query = torch.cat((query_rot, query_pass), dim=-1) + key = torch.cat((key_rot, key_pass), dim=-1) + else: + query = query_rot + key = key_rot + query = query.flatten(-2) + key = key.flatten(-2) + return query, key + + +def vllm_batched_rotary_embedding( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + is_neox_style: bool, + rotary_dim: int, + offsets: torch.Tensor, +): + + """ + Args: + positions: (num_tokens) torch.long + query: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16 + key: (num_tokens, num_heads * head_size) torch.half, torch.float, torch.bfloat16 + head_size: int + cos_sin_cache: (max_position, head_size) torch.half, torch.float, torch.bfloat16 + is_neox_style: bool + rotary_dim: int + offsets: (positions, head_size) torch.int64 + Returns: + None. 对query, key 做in place 操作 + """ + ops.infer.vllm_batched_rotary_embedding( + positions, + query, + key, + head_size, + cos_sin_cache, + is_neox_style, + rotary_dim, + offsets, + ) + + +def ref_vllm_reshape_and_cache_v4( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, +): + num_tokens, num_heads, head_size = key.shape + x = 16 // torch.tensor([], dtype=key.dtype).element_size() + block_size = key_cache.size(3) + + reshaped_key = key.reshape(num_tokens, num_heads, head_size // x, x) + for i in range(num_tokens): + block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor") + block_offset = slot_mapping[i] % block_size + key_cache[block_idx, :, :, block_offset, :] = reshaped_key[i] + value_cache[block_idx, :, :, block_offset] = value[i] + + +def vllm_reshape_and_cache_v4( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, +): + + """ + Args: + key: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + value: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + key_cache: (num_blocks, num_heads, head_size//8, block_size, 8) or (num_blocks, num_heads, head_size//4, block_size, 4) torch.half, torch.float, torch.bfloat16 + if dtype=torch.half or torch.bfloat16,key_cache shape: (num_blocks, num_heads, head_size//8, block_size, 8) + if dtype=torch.float,key_cache shape: (num_blocks, num_heads, head_size//4, block_size, 4) + value_cache: (num_blocks, num_heads, head_size//8, block_size, 8) or (num_blocks, num_heads, head_size//4, block_size, 4) torch.half, torch.float, torch.bfloat16 + if dtype=torch.half or torch.bfloat16,value_cache shape: (num_blocks, num_heads, head_size//8, block_size, 8) + if dtype=torch.float,value_cache shape: (num_blocks, num_heads, head_size//4, block_size, 4) + slot_mapping: (num_tokens) torch.long + Returns: + None, 对key_cache,value_cache进行in place 操作 + """ + ops.infer.vllm_reshape_and_cache_v4( + key, + value, + key_cache, + value_cache, + slot_mapping, + key.stride(0), + value.stride(0), + ) + + +def vllm_cache_ops_reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, +): + + """ + Args: + key: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + value: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16 + key_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.bfloat16 + value_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.bfloat16 + slot_mapping: (num_tokens) torch.long + Returns: + None, 对key_cache,value_cache进行in place 操作 + """ + + num_tokens, num_kv_heads, head_size = key.shape + num_blocks = key_cache.size(0) + key_cache = key_cache.view(num_blocks, num_kv_heads, -1, head_size) + value_cache = value_cache.view(num_blocks, num_kv_heads, -1, head_size) + ops.infer.vllm_cache_ops_reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + key.stride(0), + value.stride(0), + ) + + +def ref_vllm_reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, +): + num_tokens, _, _ = key.shape + block_size = key_cache.size(2) + v_dim = value.shape[-1] + + for i in range(num_tokens): + block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor") + block_offset = slot_mapping[i] % block_size + key_cache[block_idx, :, block_offset, :] = key[i] + value_cache[block_idx, :, block_offset, :v_dim] = value[i] + + +def vllm_reshape_and_cache( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, +): + + """ + Args: + key: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + value: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + key_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.float, torch.bfloat16 + value_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.float, torch.bfloat16 + slot_mapping: (num_tokens) torch.long + Returns: + None, 对key_cache,value_cache进行in place 操作 + """ + ops.infer.vllm_reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + key.stride(0), + value.stride(0), + ) + + +def ref_reshape_and_cache_flash( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: float, + v_scale: float, +): + num_tokens, _, _ = key.shape + block_size = key_cache.size(2) + + for i in range(num_tokens): + block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor") + block_offset = slot_mapping[i] % block_size + key_cache[block_idx, :, block_offset, :] = key[i] + value_cache[block_idx, :, block_offset, :] = value[i] + + +def reshape_and_cache_flash( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: float, + v_scale: float, +) -> None: + + """ + Args: + key: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + value: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16 + key_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.float, torch.bfloat16 + value_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.float, torch.bfloat16 + slot_mapping: (num_tokens) torch.long + kv_cache_dtype: str + k_scale: float + v_scale: float + Returns: + None, 对key_cache,value_cache进行in place 操作 + """ + assert k_scale == 1 and v_scale == 1 + assert kv_cache_dtype == "auto" + + ops.infer.vllm_reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + key.stride(0), + value.stride(0), + ) + + +def ref_vllm_copy_blocks( + key_caches, + value_caches, + block_mapping, +): + for k, v in zip(key_caches, value_caches): + src = block_mapping[:, 0] + dst = block_mapping[:, 1] + k[dst] = k[src] + v[dst] = v[src] + + +def vllm_copy_blocks( + key_caches, + value_caches, + block_mapping, +): + + """ + Args: + key_caches: [(num_blocks, num_heads, block_size, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16] + value_caches: [(num_blocks, num_heads, block_size, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16] + block_mapping: (num_tokens, 2) torch.int64 + Returns: + None, 对key_caches,value_caches进行in place 操作 + """ + ops.infer.vllm_copy_blocks( + key_caches, + value_caches, + block_mapping, + ) + + +def ref_vllm_swap_blocks(src, dst, mapping): + for item in mapping: + src_idx = item[0] + dst_idx = item[1] + dst[dst_idx] = src[src_idx].to(dst.device) + + +def vllm_swap_blocks(src: "torch.Tensor", dst: "torch.Tensor", mapping: "torch.Tensor"): + + """ + Args: + src: [(num_blocks, num_kv_heads, block_size, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16] + dst: [(num_blocks, num_kv_heads, block_size, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16] + mapping: (num_tokens, 2) torch.int64 + Returns: + None, 对dst进行in place 操作 + """ + ops.infer.vllm_swap_blocks(src, dst, mapping) + + +def ref_vllm_concat_and_cache_mla( + kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale +): + num_tokens, kv_lora_rank = kv_c.shape + _, _, pe_dim = k_pe.shape + _, block_size, rope_dim = kv_cache.shape + assert kv_lora_rank + pe_dim == rope_dim + + for i in range(num_tokens): + block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor") + block_offset = slot_mapping[i] % block_size + kv_cache[block_idx, block_offset, :kv_lora_rank] = kv_c[i] + kv_cache[block_idx, block_offset, kv_lora_rank:] = k_pe[i, 0] + + +def ref_vllm_gather_cache( + src_cache: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + dst: torch.Tensor, # [TOT_TOKENS, ENTRIES...] + block_table: torch.Tensor, # [BATCH, BLOCK_INDICES] + cu_seq_lens: torch.Tensor, # [BATCH+1] + batch_size: int, + seq_starts: torch.Tensor = None, +): + # 验证输入张量的设备一致性 + assert src_cache.device == dst.device == block_table.device == cu_seq_lens.device + if seq_starts is not None: + assert seq_starts.device == src_cache.device + + # 获取基本参数 + block_size = src_cache.size(1) + entry_size = src_cache.flatten(2, -1).size(2) + + # 处理每个批次 + for bid in range(batch_size): + seq_start = cu_seq_lens[bid] + seq_end = cu_seq_lens[bid + 1] + seq_len = seq_end - seq_start + + # 计算需要的块数 + tot_blocks = math.ceil(seq_len / block_size) + + # 获取当前批次的块表 + if seq_starts is not None: + offset = seq_starts[bid] // block_size + batch_block_table = block_table[bid, offset : offset + tot_blocks] + else: + batch_block_table = block_table[bid, :tot_blocks] + + # 准备目标位置 + dst_seq = dst[seq_start:seq_end] + + # 处理完整块 + full_blocks = seq_len // block_size + if full_blocks > 0: + # 获取所有完整块的源数据 [full_blocks, block_size, entry_size] + src_blocks = src_cache[batch_block_table[:full_blocks]] + # 展平并复制到目标位置 + dst_seq[: full_blocks * block_size].copy_(src_blocks.flatten(0, 1)) + + # 处理部分块 + partial_size = seq_len % block_size + if partial_size > 0: + last_block = src_cache[batch_block_table[full_blocks], :partial_size] + dst_seq[full_blocks * block_size :].copy_(last_block) + + +def vllm_gather_cache( + src_cache: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + dst: torch.Tensor, # [TOT_TOKENS, ENTRIES...] + block_table: torch.Tensor, # [BATCH, BLOCK_INDICES] + cu_seq_lens: torch.Tensor, # [BATCH+1] + batch_size: int, + seq_starts: torch.Tensor = None, +): + """ + Args: + src_cache: [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] torch.float16, torch.bfloat16 int + dst: [TOT_TOKENS, ENTRIES...] torch.float16, torch.bfloat16 + block_table: [BATCH, BLOCK_INDICES] torch.int + cu_seq_lens: [BATCH+1] torch.int + batch_size: int + seq_starts: [BATCH] or None torch.int + """ + ops.infer.vllm_gather_cache( + src_cache, dst, block_table, cu_seq_lens, batch_size, seq_starts + ) + + +def ref_vllm_gather_cache_int8( + src_cache: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + src_cache_scale: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, 2] + kv_lora_rank: int, + dst: torch.Tensor, # [TOT_TOKENS, ENTRIES...] + block_table: torch.Tensor, # [BATCH, BLOCK_INDICES] + cu_seq_lens: torch.Tensor, # [BATCH+1] + batch_size: int, + seq_starts: torch.Tensor = None, +): + # 验证输入张量的设备一致性 + assert ( + src_cache.device + == src_cache_scale.device + == dst.device + == block_table.device + == cu_seq_lens.device + ) + if seq_starts is not None: + assert seq_starts.device == src_cache.device + + # 获取基本参数 + block_size = src_cache.size(1) + + # 处理每个批次 + for bid in range(batch_size): + seq_start = cu_seq_lens[bid] + seq_end = cu_seq_lens[bid + 1] + seq_len = seq_end - seq_start + + # 计算需要的块数 + tot_blocks = math.ceil(seq_len / block_size) + + # 获取当前批次的块表 + if seq_starts is not None: + offset = seq_starts[bid] // block_size + batch_block_table = block_table[bid, offset : offset + tot_blocks] + else: + batch_block_table = block_table[bid, :tot_blocks] + + # 准备目标位置 + dst_seq = dst[seq_start:seq_end] + + # 处理完整块 + full_blocks = seq_len // block_size + if full_blocks > 0: + # 获取所有完整块的源数据 [full_blocks, block_size, entry_size] + src_cache_blocks = src_cache[batch_block_table[:full_blocks]] + src_scale_blocks = src_cache_scale[batch_block_table[:full_blocks]] + src_k_cache_blocks = src_cache_blocks[ + ..., :kv_lora_rank + ] * src_scale_blocks[..., 0].unsqueeze(-1) + src_k_pe_blocks = src_cache_blocks[..., kv_lora_rank:] * src_scale_blocks[ + ..., 1 + ].unsqueeze(-1) + src_blocks = torch.cat((src_k_cache_blocks, src_k_pe_blocks), dim=-1).to( + dst.dtype + ) + # 展平并复制到目标位置 + dst_seq[: full_blocks * block_size].copy_(src_blocks.flatten(0, 1)) + + # 处理部分块 + partial_size = seq_len % block_size + if partial_size > 0: + last_block = src_cache[batch_block_table[full_blocks], :partial_size] + last_src_scale_blocks = src_cache_scale[ + batch_block_table[full_blocks], :partial_size + ] + last_src_k_cache_blocks = last_block[ + ..., :kv_lora_rank + ] * last_src_scale_blocks[..., 0].unsqueeze(-1) + last_src_k_pe_blocks = last_block[ + ..., kv_lora_rank: + ] * last_src_scale_blocks[..., 1].unsqueeze(-1) + last_block = torch.cat( + (last_src_k_cache_blocks, last_src_k_pe_blocks), dim=-1 + ).to(dst.dtype) + dst_seq[full_blocks * block_size :].copy_(last_block) + + +def vllm_gather_cache_int8( + src_cache: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] + src_cache_scale: torch.Tensor, # [NUM_BLOCKS, BLOCK_SIZE, 2] + kv_lora_rank: int, + dst: torch.Tensor, # [TOT_TOKENS, ENTRIES...] + block_table: torch.Tensor, # [BATCH, BLOCK_INDICES] + cu_seq_lens: torch.Tensor, # [BATCH+1] + batch_size: int, + seq_starts: torch.Tensor = None, +): + """ + Args: + src_cache: [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...] torch.int8 + src_cache_scale: [NUM_BLOCKS, BLOCK_SIZE, 2] torch.float32 + kv_lora_rank: int + dst: [TOT_TOKENS, ENTRIES...] torch.float16, torch.bfloat16 + block_table: [BATCH, BLOCK_INDICES] torch.int + cu_seq_lens: [BATCH+1] torch.int + batch_size: int + seq_starts: [BATCH] or None torch.int + """ + ops.infer.vllm_gather_cache_int8( + src_cache, + src_cache_scale, + kv_lora_rank, + dst, + block_table, + cu_seq_lens, + batch_size, + seq_starts, + ) + + +def ref_vllm_concat_and_cache_mla_int8( + kv_c_int8: torch.Tensor, + kv_c_scale: torch.Tensor, + k_pe_int8: torch.Tensor, + k_pe_scale: torch.Tensor, + kv_cache: torch.Tensor, + kv_cache_scale: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + scale: torch.Tensor, +) -> None: + + num_tokens, kv_lora_rank = kv_c_int8.shape + _, block_size, _ = kv_cache_scale.shape + + for i in range(num_tokens): + block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor") + block_offset = slot_mapping[i] % block_size + kv_cache[block_idx, block_offset, :kv_lora_rank] = kv_c_int8[i] + kv_cache[block_idx, block_offset, kv_lora_rank:] = k_pe_int8[i][0] + kv_cache_scale[block_idx, block_offset, 0] = kv_c_scale[i] + kv_cache_scale[block_idx, block_offset, 1] = k_pe_scale[i][0] + + +def vllm_concat_and_cache_mla( + kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale +): + ops.infer.vllm_concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping) + + +def vllm_concat_and_cache_mla_int8( + kv_c_int8, + kv_c_scale, + k_pe_int8, + k_pe_scale, + kv_cache, + kv_cache_scale, + slot_mapping, + kv_cache_dtype, + scale, +): + """ + Args: + kv_c_int8: [num_tokens, kv_lora_rank] torch.int8 + kv_c_scale: [num_tokens] torch.float32 + k_pe_int8: [num_tokens, n, pe_dim] torch.int8 + k_pe_scale: [num_tokens, n] torch.float32 + kv_cache: [num_blocks, block_size, (kv_lora_rank + pe_dim)] torch.int8 + kv_cache_scale: [num_blocks, block_size, 2] torch.float32 + slot_mapping: [num_tokens] torch.long + """ + + ops.infer.vllm_concat_and_cache_mla_int8( + kv_c_int8, + kv_c_scale, + k_pe_int8, + k_pe_scale, + kv_cache, + kv_cache_scale, + slot_mapping, + ) + + +class vllm_llama_mlp(CF.VllmLlamaMlp): + def __init__( + self, + gate_up_proj_weight: "torch.Tensor", + down_proj_weight: "torch.Tensor", + hidden_size: int, + intermediate_size: int, + tp: int, + ) -> None: + gate_up_proj_weight = gate_up_proj_weight + down_proj_weight = down_proj_weight + hidden_size = hidden_size + intermediate_size = intermediate_size + tp = tp + super().__init__( + gate_up_proj_weight, + down_proj_weight, + hidden_size, + intermediate_size, + tp, + ) + + def __call__(self, x: "torch.Tensor", group=None): + x1 = x + if group is None: + super().forward(x1, x1) + else: + from ixformer.distributed._distributed import _check_group + + group = _check_group(group) + super().forward(x1, x1, group) + return x + + +def gptq_gemm( + input: torch.Tensor, + qweight: torch.Tensor, + qzeros: torch.Tensor, + scales: torch.Tensor, + g_idx: torch.Tensor, + use_exllama: bool, + weight_bits: int, +) -> torch.Tensor: + """ + use_exllama + - True uesExllama + - False GeneralGptq + g_idx.is_empty() + - True don't use g_idx + - False use g_idx + 1. use_exllama == False && use g_idx + [General gptq] desc_act == True && parallel in k dimension && group_size != -1 + 2. use_exllama == True && use g_idx (g_idx has been argsort) + [Exllama with g_idx] desc_act == True && parallel in n dimension && group_size != -1 + 3. use_exllama == True && don't use g_idx + [Exllama] desc_act == False || desc_act == True && group_size == -1 + + Args: + input: (m, k) torch.float16, torch.bfloat16 + qweight: (k // (32 / bits), n) torch.int32 + qzeros: (k / group_size, n / (32 / bits)) torch.int32 + scales: (k // group_size, n) torch.float16, torch.bfloat16 + g_idx: (k) torch.int32 + use_exllama: bool + wheather use exllama + weight_bits: int + quant bits of weight + Returns: + output: (m, n) torch.float16, torch.bfloat16 + """ + bs = input.shape[0] + group_size = input.shape[1] // scales.shape[0] + + # condition : without gidx or group_size == -1 + ixinfer_gemm_supported = ( + weight_bits == 4 + and (g_idx is None or g_idx.numel() == 0) + and (scales.shape[0] == 1 or group_size in [32, 128]) + ) + if use_exllama: + if bs <= 8 or ixinfer_gemm_supported: + output = ops.infer.quantized_linear( + input, + qweight, + scales, + "gptq-ex", + weight_bits, + qzeros, + None, + group_size, + g_idx, + "unknown", + ) + else: + # GPTQ GEMM TODO + o_dtype_str = "fp16" if input.dtype == torch.half else "bf16" + deq_w = ops.infer.quantized_weight_dequant( + qweight, + scales, + "gptq-ex", + o_dtype_str, + weight_bits, + qzeros, + group_size, + g_idx, + ) + + output = linear(input, deq_w.transpose(0, 1).contiguous()) + else: + if bs <= 8: + output = ops.infer.quantized_linear( + input, + qweight, + scales, + "gptq", + weight_bits, + qzeros, + None, + group_size, + g_idx, + "unknown", + ) + else: + # GPTQ GEMM TODO + o_dtype_str = "fp16" if input.dtype == torch.half else "bf16" + deq_w = ops.infer.quantized_weight_dequant( + qweight, + scales, + "gptq", + o_dtype_str, + weight_bits, + qzeros, + group_size, + g_idx, + ) + output = linear(input, deq_w.transpose(0, 1).contiguous()) + return output + + +def vllm_gptq_shuffle(qweights, g_idx, weight_bits): + ops.infer.vllm_gptq_shuffle(qweights, g_idx, weight_bits) + + +def vllm_moe_topk_softmax( + topk_weights: "torch.Tensor", + topk_ids: "torch.Tensor", + token_expert_indicies: "torch.Tensor", + gating_output: "torch.Tensor", +): + + """ + Args: + topk_weights: (num_tokens,topk) torch.float + topk_ids: (num_tokens,topk) torch.int + token_expert_indicies: (num_tokens,topk) torch.int + gating_output: (num_tokens,num_experts) torch.float + Returns: + None, 对topk_weights,topk_ids进行in place 操作 + """ + assert isinstance(topk_weights, torch.Tensor) + assert gating_output.dtype == torch.float32 + ops.infer.moe_topk_softmax( + topk_weights, topk_ids, token_expert_indicies, gating_output, False + ) + + +def vllm_moe_align_block_size( + topk_ids: "torch.Tensor", + num_experts: int, + block_size: int, + sorted_ids: "torch.Tensor", + expert_ids: "torch.Tensor", + num_tokens_post_pad: "torch.Tensor", +): + """ + Args: + topk_ids: (num_tokens,topk) torch.int + num_experts: int + block_size: int + sorted_ids: (topk_ids.numel() + num_experts * (block_size - 1)) torch.int + expert_ids: (topk_ids.numel() + num_experts) torch.int + num_tokens_post_pad: (1) torch.int + Returns: + None + """ + + ops.infer.moe_align_block_size( + topk_ids, num_experts, block_size, sorted_ids, expert_ids, num_tokens_post_pad + ) + + +def ref_vllm_invoke_fused_moe_kernel( + A: "torch.Tensor", + B: "torch.Tensor", + C: "torch.Tensor", + topk_weight: "torch.Tensor", + topk_ids: "torch.Tensor", + sorted_token_ids: "torch.Tensor", + expert_ids: "torch.Tensor", + num_tokens_post_padded: "torch.Tensor", + mul_routed_weight: bool, + top_k: int, + block_size_m: int, + persistent: bool = False, + w_scale: torch.Tensor = None, + a_scale: torch.Tensor = None, +): + + expert_num, N, K = B.shape + M, topk = C.shape[:2] + + clone_A = A.clone() + clone_B = B.clone() + if clone_A.shape[0] == M: + clone_A = clone_A.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K) + topk_ids = topk_ids.view(-1) + + if A.dtype == torch.int8: + use_scale = True + clone_A = clone_A.to(torch.float32) + clone_B = clone_B.to(torch.float32) + tmp = torch.zeros(M * topk, N, dtype=torch.float32, device=C.device) + else: + use_scale = False + tmp = torch.zeros(M * topk, N, dtype=C.dtype, device=C.device) + + for i in range(expert_num): # expert_num + mask = topk_ids == i + if mask.sum(): + tmp[mask] = clone_A[mask] @ clone_B[i].transpose(0, 1) + if use_scale: + tmp[mask] = tmp[mask] * w_scale[i].view(1, N) + + if mul_routed_weight: + tmp = tmp * topk_weight.view(-1, 1) + + if use_scale: + tmp = tmp.view(M, topk, N) + tmp = tmp * a_scale.view(M, -1, 1) + + C[:] = tmp.to(C.dtype).view(M, topk, N) + return C + + +def vllm_invoke_fused_moe_kernel( + A: "torch.Tensor", + B: "torch.Tensor", + C: "torch.Tensor", + topk_weight: "torch.Tensor", + topk_ids: "torch.Tensor", + sorted_token_ids: "torch.Tensor", + expert_ids: "torch.Tensor", + num_tokens_post_padded: "torch.Tensor", + mul_routed_weight: bool, + top_k: int, + block_size_m: int, + persistent: bool = False, + w_scale: torch.Tensor = None, + a_scale: torch.Tensor = None, +): + + """ + Args: + A: (bs*seq, K) / (bs*seq*top_k, K) torch.float16, torch.bfloat16 + B: (num_experts, N, K) torch.float16, torch.bfloat16 + C: (bs*seq, top_k, N) torch.half,torch.bfloat16 + topk_weight: (bs*seq, topk) torch.float32 + topk_ids: (bs*seq, topk) torch.int32 + sorted_token_ids: (topk_ids.numel() + num_experts * (block_size - 1)) torch.int32 + expert_ids: (topk_ids.numel() + num_experts) torch.int32 + num_tokens_post_pad:(1) torch.int32 + mul_routed_weight: bool + top_k: int + block_size_m: int + Returns: + C: (bs*seq, top_k, N) torch.half,torch.bfloat16 + """ + ops.infer.invoke_fused_moe_kernel( + A, + B, + C, + topk_weight, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + mul_routed_weight, + top_k, + block_size_m, + persistent, + w_scale, + a_scale.view(-1) * topk_weight.view(-1) + if a_scale is not None and mul_routed_weight + else a_scale, + ) + + +def advance_step_flashattn( + num_seqs: int, + num_queries: int, + block_size: int, + input_tokens: "torch.Tensor", + sampled_token_ids: "torch.Tensor", + input_positions: "torch.Tensor", + seq_lens: "torch.Tensor", + slot_mapping: "torch.Tensor", + block_tables: "torch.Tensor", +): + ops.infer.vllm_advance_step_flashattn( + num_seqs, + num_queries, + block_size, + input_tokens, + sampled_token_ids, + input_positions, + seq_lens, + slot_mapping, + block_tables, + ) + + +if config.IXFORMER_PAGED_ATTENTION_ALGO == "ixformer": + print("set IXFORMER_PAGED_ATTENTION_ALGO: ixformer") + vllm_paged_attention = vllm_paged_attention_ixformer +else: + vllm_paged_attention = vllm_paged_attention_ixinfer diff --git a/ixformer_sdk/inference/functions/w8a16.py b/ixformer_sdk/inference/functions/w8a16.py new file mode 100644 index 00000000..4071b657 --- /dev/null +++ b/ixformer_sdk/inference/functions/w8a16.py @@ -0,0 +1,225 @@ +import math +from typing import List, Union, Optional + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +import ixformer +from ixformer.core import config + +__all__ = [ + "w8a16_gemm", + "w8a16_gemv", + "w8a16", + "ref_w8a16", + "wu8a16", + "ref_wu8a16", +] + + +def w8a16_gemv( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + group_size: int = -1, + format: str = "unknown", + output: Optional[torch.Tensor] = None +): + """ + w8a16 gemv 接口 + input : bf16|fp16 (bs, ic) + qweights : int8 TN:(oc, ic) NN:(ic, oc) + scales : bf16|fp16 TN: 当groupsize为-1时, shape: (1, oc), 否则,shape: (ic/group_size, oc) NN:(1, oc) + TN 支持条件： ic % groupSize = 0, oc % 2 = 0, bs<=4 + NN 支持条件： groupsize = -1 or groupsize = ic, oc % 4 = 0, bs<=4 + """ + assert format in ["TN", "NN"] + assert len(qweights.shape) == 2 + assert len(scales.shape) == 2 + + input_shape = list(inputs.shape) + inputs = inputs.view(-1, input_shape[-1]) + + if format == "TN": + output_shape = input_shape[:-1] + [qweights.shape[0]] + else: + output_shape = input_shape[:-1] + [qweights.shape[1]] + + if output is None: + output = inputs.new_empty(output_shape).view(-1, output_shape[-1]) + + ops.infer.w8a16_gemv(output, inputs, qweights, scales, group_size, format) + return output.view(output_shape) + + +def w8a16_gemm( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + group_size: int = -1, + format: str = "TN", + persistent: int = 0, + output: Optional[torch.Tensor] = None +): + """ + w8a16 gemm 接口 + 1. group_size=-1 or group_size=ic + input : bf16|fp16 (bs, ic) + qweights : int8 TN:(oc, ic) NN:(ic, oc) + scales : bf16|fp16 (1, oc) + NN 支持条件： ic%64==0, oc%64==0 + + 2. group_size=64 + input : bf16|fp16 (bs, ic) + qweights : int8 TN:(oc, ic) + scales : bf16|fp16 (ic/64, oc) + TN 支持条件： oc%2==0, ic%64==0 + NN 不支持 + """ + + assert format in ["TN", "NN"] + assert len(qweights.shape) == 2 + assert len(scales.shape) == 2 + + input_shape = list(inputs.shape) + inputs = inputs.view(-1, input_shape[-1]) + + if format == "TN": + output_shape = input_shape[:-1] + [qweights.shape[0]] + else: + output_shape = input_shape[:-1] + [qweights.shape[1]] + + if output is None: + output = inputs.new_empty(output_shape).view(-1, output_shape[-1]) + + ops.infer.w8a16_gemm( + output, inputs, qweights, scales, group_size, format, persistent + ) + return output.view(output_shape) + + +def dequant(qweight, scales, group_size): + IC, OC = qweight.shape + weight = qweight.t().reshape(OC, -1, group_size).to( + torch.float32 + ) * scales.t().unsqueeze(-1) + return weight.reshape(OC, IC) + + +def ref_w8a16( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + group_size: int = -1, + format: str = "TN", +): + if group_size == -1: + group_size = inputs.shape[1] + if format == "TN": + weights = dequant(qweights.transpose(0, 1), scales, group_size) + elif format == "NN": + weights = dequant(qweights, scales, group_size) + return torch.nn.functional.linear(inputs, weights.to(inputs.dtype)) + + +def w8a16( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + group_size: int = -1, + format: str = "TN", + output: Optional[torch.Tensor] = None, + persistent: int = 0, +): + input_shape = inputs.shape + inputs = inputs.view(-1, input_shape[-1]) + bs = inputs.size(0) + inputs = inputs.view(input_shape) + if bs <= config.IXFORMER_GEMV_THRESHOLD: + return w8a16_gemv( + inputs=inputs, + qweights=qweights, + scales=scales, + group_size=group_size, + format=format, + output=output + ) + else: + return w8a16_gemm( + inputs=inputs, + qweights=qweights, + scales=scales, + group_size=group_size, + format=format, + output=output, + persistent=persistent + ) + + +def ref_wu8a16( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + group_size: int = -1, + format: str = "TN", +): + assert format in ["TN"] + assert len(qweights.shape) == 2 + assert len(scales.shape) == 2 + org_w_shape = qweights.shape + scales = scales.transpose(0, 1).flatten().view(-1, 1) + zeros = zeros.transpose(0, 1).flatten().view(-1, 1) + + if group_size != -1: + qweights = qweights.reshape(-1, group_size) + w = (qweights - zeros) * scales + w = w.reshape(org_w_shape) + output = torch.matmul(inputs, w.t()) + return output + + +def wu8a16( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + group_size: int = -1, + format: str = "TN", + persistent: int = 0, +): + """ + http://confluence.iluvatar.ai:8090/display/SW/cuinferCustomGemm+Interface+Doc + wu8a16 非对称量化 gemm 接口 + 1. group_size=-1 + input : bf16|fp16 (bs, ic) + qweights : uint8 TN:(oc, ic) + scales : bf16|fp16 (1,oc) + zeros : bf16|fp16 (1,oc) + TN 支持条件： ic % 64 == 0 + NN 不支持 + + 2. group_size=64 + input : bf16|fp16 (bs, ic) + qweights : int8 TN:(oc, ic) + scales : bf16|fp16 (ic/64, oc) + zeros : bf16|fp16 (ic/64, oc) + TN 支持条件： oc % 2 == 0 && ic % 64 == 0 + NN 不支持 + """ + assert format in ["TN"] + assert len(qweights.shape) == 2 + assert len(scales.shape) == 2 + + input_shape = list(inputs.shape) + inputs = inputs.view(-1, input_shape[-1]) + + output_shape = input_shape[:-1] + [qweights.shape[0]] + + output = inputs.new_empty(output_shape).view(-1, output_shape[-1]) + + ops.infer.wu8a16_gemm( + output, inputs, qweights, scales, zeros, group_size, format, persistent + ) + return output.view(output_shape) diff --git a/ixformer_sdk/inference/functions/w8a8.py b/ixformer_sdk/inference/functions/w8a8.py new file mode 100644 index 00000000..02434cbc --- /dev/null +++ b/ixformer_sdk/inference/functions/w8a8.py @@ -0,0 +1,317 @@ +from typing import Optional, Tuple + +import ixformer._C as ops +import torch + +__all__ = [ + "w8a8", + "ref_w8a8", + "dynamic_scaled_int8_quant", + "ref_dynamic_scaled_int8_quant", + "static_scaled_int8_quant", + "ref_static_scaled_int8_quant", + "scaled_int8_quant", +] + + +def ref_w8a8( + input: "torch.Tensor", + weight: "torch.Tensor", + i_scales: "torch.Tensor", + w_scales: "torch.Tensor", + output: "torch.Tensor", + format: str = "TN", + persistent=0, + bias: torch.Tensor = None, +): + dtype = output.dtype + input_f32 = input.to(torch.float32) + weight_f32 = weight.to(torch.float32) + assert format in ["TN", "NN", "NT"] + if format == "TN": + weight_f32 = weight_f32.transpose(0, 1) + if format == "NT": + input_f32 = input_f32.transpose(0, 1) + output_f32 = ( + torch.matmul(input_f32, weight_f32) + * i_scales.view(-1, 1) + * w_scales.view(1, -1) + ) + if bias is not None: + bias_f32 = bias.to(torch.float32) + output_f32 += bias_f32.view(1, -1) + output.copy_(output_f32.to(dtype)) + return output + + +def w8a8_gemm( + input: "torch.Tensor", + weight: "torch.Tensor", + i_scales: "torch.Tensor", + w_scales: "torch.Tensor", + bias: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + format: str = "TN", + persistent: bool = False, + out_dtype: torch.dtype = None, +): + """ + Args: + input: (n, k) torch.int8 + weight: (m, k) if format == "TN" else (k, m) torch.int8 + i_scales: (n) torch.float32 + w_scales: (m) torch.float32 + bias: (m) torch.float32, same as output_type + format: str + Options include TN, NN and NT + persistent: Whether to use overleap bool + out_dtype: torch.float16, torch.bfloat16 + Returns: + output: (n, m) torch.float16, torch.bfloat16 + """ + + input_shape = input.shape + + if output is None: + if out_dtype is None: + raise RuntimeError("w8a8 gemm need out_dtype argument when output is none.") + output = torch.empty( + (input_shape[:-1] + (weight.shape[0],)), + dtype=out_dtype, + device=input.device, + ) + + output_shape = output.shape + + input = input.view(-1, input_shape[-1]) + output = output.view(-1, output_shape[-1]) + + ops.infer.w8a8_gemm( + output, input, weight, i_scales, w_scales, bias, format, int(persistent) + ) + + return output.view(*output_shape) + + +def ref_static_scaled_int8_quant(output, input, scale): + """ + Args: + output: [torch.int8] [m, k] + input: [torch.half,torch.bfloat16] [m, k] + scale: [torch.float32] [1] + Returns: + output: [torch.int8] [m, k] + scale: [torch.float32] [1] + """ + # [m, 1] + f_input = input / scale.to(input.dtype) + i_output = torch.clamp(torch.round(f_input), -127, 127).to(torch.int8) + output.copy_(i_output) + return output, scale + + +# for vllm: https://github.com/vllm-project/vllm/blob/v0.5.4/vllm/_custom_ops.py#L387 +def static_scaled_int8_quant(output, input, scale): + """ + Args: + output: [torch.int8] [m, k] + input: [torch.half,torch.bfloat16] [m, k] + scale: [torch.float32] [1] + Returns: + output: [torch.int8] [m, k] + scale: [torch.float32] [1] + """ + ops.infer.scaled_int8_quant(output, input, scale, 0) + return output, scale + + +def ref_dynamic_scaled_int8_quant(output, input, scale): + """ + Args: + output: [torch.int8] [m, k] + input: [torch.half,torch.bfloat16] [m, k] + scale: [torch.float32] [m] + Returns: + output: [torch.int8] [m, k] + scale: [torch.float32] [m] + """ + # [m, 1] + amax_, _ = torch.max(torch.abs(input), dim=-1, keepdim=True) + f_scale = amax_.float() / 127.0 + scale.view(-1).copy_(f_scale.view(-1)) + + f_input = input / f_scale.to(input.dtype) + i_output = torch.clamp(torch.round(f_input), -127, 127).to(torch.int8) + output.copy_(i_output) + return output, scale.view(input.shape[:-1]) + + +# for vllm: https://github.com/vllm-project/vllm/blob/v0.5.4/vllm/_custom_ops.py#L394 +def dynamic_scaled_int8_quant(output, input, scale): + """ + Args: + output: [torch.int8] [m, k] + input: [torch.half,torch.bfloat16] [m, k] + scale: [torch.float32] [m] + Returns: + output: [torch.int8] [m, k] + scale: [torch.float32] [m] + """ + ops.infer.scaled_int8_quant(output, input, scale, 1) + return output, scale + + +def scaled_int8_quant( + input: torch.Tensor, scale: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Quantize the input tensor to int8 and return the quantized tensor and scale. + + Args: + input: The input tensor to be quantized to int8. + scale: Optional scaling factor for the int8 quantization. + When not provided, we invoke dynamic-per-token quantization. + + Returns: + Tuple[Torch.Tensor, Torch.Tensor] : Output int8 tensor and scales. + """ + output = torch.empty_like(input, dtype=torch.int8) + if scale is not None: + # static-per-tensor quantization. + static_scaled_int8_quant(output, input, scale) + return output, scale + + # dynamic-per-token quantization. + input_scales = torch.empty( + (input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32 + ) + dynamic_scaled_int8_quant(output, input, input_scales) + return output, input_scales + + +def w8a8_gemv( + input: "torch.Tensor", + weight: "torch.Tensor", + i_scales: "torch.Tensor", + w_scales: "torch.Tensor", + bias: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + format: str = "TN", + persistent: bool = False, + out_dtype: torch.dtype = None, +): + """ + Args: + input: (n, k) torch.int8 + weight: (m, k) if format == "TN" else (k, m) torch.int8 + i_scales: (n) torch.float32 + w_scales: (m) torch.float32 + bias: (m) torch.float32 same as output_type + format: str + Options include TN and NN + persistent: Whether to use overleap bool + out_dtype: torch.float16, torch.bfloat16 + Returns: + output: (n, m) torch.float16, torch.bfloat16 + """ + + input_shape = input.shape + + if output is None: + if out_dtype is None: + raise RuntimeError("w8a8 gemv need out_dtype argument when output is none.") + output = torch.empty( + (input_shape[:-1] + (weight.shape[0],)), + dtype=out_dtype, + device=input.device, + ) + + input = input.view(-1, input_shape[-1]) + + ops.infer.w8a8_gemv( + output, input, weight, i_scales, w_scales, bias, format, int(persistent) + ) + + return output + + +def handle_pading(weight: torch.Tensor, format: str, is_gemm: bool): + """Handle padding alignment for weight matrices + Args: + weight: Original weight matrix [m, k] + format: Matrix format, TN indicates transposed layout + is_gemm: Whether for GEMM operation (requires extra alignment checks) + Returns: + torch.Tensor: Padded weight matrix + Raises: + AssertionError: When is_gemm=True requires 4-byte alignment for m/k + """ + # weight should have been pad before w8a8 is called, handle _padding here just ensure the code run success, + # but performance is low, please refer to vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py + m, k = weight.shape + s = weight.stride(0) + if s % 64 != 0 and format == "TN": + pad_k = (s // 64 + 1) * 64 + weight_pad = torch.empty((m, pad_k), dtype=weight.dtype, device=weight.device) + _weight = weight_pad[:, :k] + if is_gemm: + assert m % 4 == 0 and k % 4 == 0 + _weight.copy_(weight) + return _weight + else: + return weight + + +def w8a8( + input: torch.Tensor, + weight: torch.Tensor, + i_scales: torch.Tensor, + w_scales: torch.Tensor, + bias: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + format: str = "TN", + persistent: bool = False, + out_dtype: torch.dtype = None, +): + """ + Args: + input: (n, k) torch.int8 + weight: (m, k) if format == "TN" else (k, m) torch.int8 + i_scales: (n) torch.float32 + w_scales: (m) torch.float32 + bias: (m) torch.float32, same as output_type + format: str + Options include TN and NN + persistent: Whether to use overleap bool + out_dtype: torch.float16, torch.bfloat16 + Returns: + output: (n, m) torch.float16, torch.bfloat16 + """ + bs = input.numel() // input.shape[-1] + gemv_condition = (format == "TN" and bs <= 1) or (format == "NN" and bs <= 16) + if gemv_condition: + weight = handle_pading(weight, format, is_gemm=False) + return w8a8_gemv( + input, + weight, + i_scales, + w_scales, + bias=bias, + output=output, + format=format, + persistent=persistent, + out_dtype=out_dtype, + ) + else: + weight = handle_pading(weight, format, is_gemm=True) + return w8a8_gemm( + input, + weight, + i_scales, + w_scales, + bias=bias, + output=output, + format=format, + persistent=persistent, + out_dtype=out_dtype, + ) diff --git a/ixformer_sdk/inference/functions/wi4a16.py b/ixformer_sdk/inference/functions/wi4a16.py new file mode 100644 index 00000000..413b21e0 --- /dev/null +++ b/ixformer_sdk/inference/functions/wi4a16.py @@ -0,0 +1,157 @@ +import ixformer._C as ops +import torch + +__all__ = ["wi4a16_gemm", "wi4a16_gemv", "wi4a16", "ref_wi4a16"] + + +def dequant_weight(tensor, scales, zeros, block_size): + # from CPM + """ + tensor: (oc/2, ic) + scales: (oc, ic/group_size) + zeros: (oc, ic/group_size) + """ + dtype = scales.dtype + left = tensor >> 4 + right = tensor << 4 >> 4 + left, right = right, left + ret = torch.cat((left, right), dim=-1).reshape(-1, left.size(-1)) + ret_shape = ret.size() + ret = ret.view(-1, block_size) + ret = scales.view(-1, 1) * (ret - zeros.view(-1, 1)) + ret = ret.reshape(ret_shape).to(dtype=dtype) + return ret + + +def ref_wi4a16( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + group_size: int = -1, + format: str = "TN", +): + assert format in ["TN"] + weights = dequant_weight( + qweights, + scales.transpose(0, 1).contiguous(), + zeros.transpose(0, 1).contiguous(), + group_size, + ) + output = torch.nn.functional.linear(inputs, weights.to(inputs.dtype)) + return output + + +def wi4a16_gemm( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + group_size: int = -1, + format: str = "TN", + output=None, +): + """ + wi4a16 gemm 接口 + 支持条件: + format = TN + group_size = 128 + input : fp16 (bs, ic) + qweights : int8 (oc/2, ic) + scales : fp16 (ic/group_size, oc) + zeros : fp16 (ic/group_size, oc) + TN 支持条件： oc % 2 == 0 && ic % 128 == 0 + NN 支持条件： 不支持 + """ + + assert format in ["TN"] + assert len(qweights.shape) == 2 + assert len(scales.shape) == 2 + assert len(zeros.shape) == 2 + + input_shape = list(inputs.shape) + inputs = inputs.view(-1, input_shape[-1]) + + if output is None: + output_shape = input_shape[:-1] + [scales.shape[1]] + output = inputs.new_empty(output_shape).view(-1, output_shape[-1]) + else: + output_shape = output.shape + + ops.infer.wi4a16_gemm(output, inputs, qweights, scales, zeros, group_size, format) + return output.view(output_shape) + + +def wi4a16_gemv( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + group_size: int = -1, + format: str = "TN", + output=None, +): + """ + wi4a16 gemv 接口 + 支持条件: + format = TN + group_size = 128 + input : bf16|fp16 (bs, ic) + qweights : int8 (oc/2, ic) + scales : bf16|fp16 (ic/group_size, oc) + zeros : bf16|fp16 (ic/group_size, oc) + TN 支持条件： oc % 2 == 0 && ic % 128 == 0 + NN 支持条件： 不支持 + """ + + assert format in ["TN"] + assert len(qweights.shape) == 2 + assert len(scales.shape) == 2 + assert len(zeros.shape) == 2 + + input_shape = list(inputs.shape) + inputs = inputs.view(-1, input_shape[-1]) + + if output is None: + output_shape = input_shape[:-1] + [scales.shape[1]] + output = inputs.new_empty(output_shape).view(-1, output_shape[-1]) + else: + output_shape = output.shape + + ops.infer.wi4a16_gemv(output, inputs, qweights, scales, zeros, group_size, format) + return output.view(output_shape) + + +def wi4a16( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + group_size: int = -1, + format: str = "TN", + output=None, +): + input_shape = inputs.shape + inputs = inputs.view(-1, input_shape[-1]) + bs = inputs.size(0) + inputs = inputs.view(input_shape) + if bs <= 1: + return wi4a16_gemv( + inputs=inputs, + qweights=qweights, + scales=scales, + zeros=zeros, + group_size=group_size, + format=format, + output=output, + ) + else: + return wi4a16_gemm( + inputs=inputs, + qweights=qweights, + scales=scales, + zeros=zeros, + group_size=group_size, + format=format, + output=output, + ) diff --git a/ixformer_sdk/inference/functions/wui4a16.py b/ixformer_sdk/inference/functions/wui4a16.py new file mode 100644 index 00000000..243b0284 --- /dev/null +++ b/ixformer_sdk/inference/functions/wui4a16.py @@ -0,0 +1,155 @@ +import ixformer._C as ops +import torch + +__all__ = ["wui4a16_gemm", "wui4a16_gemv", "wui4a16", "ref_wui4a16"] + + +def dequant_weight(tensor, scales, zeros, block_size): + """ + tensor: (oc/2, ic) + scales: (oc, ic/group_size) + zeros: (oc, ic/group_size) + """ + dtype = scales.dtype + left = tensor >> 4 + right = tensor << 4 >> 4 + left, right = right, left + ret = torch.cat((left, right), dim=-1).reshape(-1, left.size(-1)) + ret_shape = ret.size() + ret = ret.view(-1, block_size) + ret = scales.view(-1, 1) * (ret - zeros.view(-1, 1)) + ret = ret.reshape(ret_shape).to(dtype=dtype) + return ret + + +def ref_wui4a16( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + bias: "torch.Tensor" = None, + group_size: int = -1, + format: str = "NN", + only_return_weight: bool = False, +): + """ + format = TN,NN + group_size = TN(128),NN(128, 32) + input : bfloat16|fp16 (bs, ic) + qweights : int32 NN: (ic, oc // 8) TN:(oc, ic // 8) + scales : bfloat16|fp16 (ic // group_size, oc) + zeros : int32 (ic // group_size, oc // 8) + bias : bfloat16|fp16 (oc, ) + output : bfloat16|fp16 (bs, oc) + """ + + def unpack_tensor(x, pack_num=8, order_map=None): + if order_map is None: + order_map = [0, 1, 2, 3, 4, 5, 6, 7] + unit = 32 // pack_num + rows, cols = x.shape + res = torch.zeros((rows, cols * pack_num), dtype=torch.int32, device=x.device) + for col in range(cols): + for k in range(pack_num): + res[:, col * pack_num + order_map[k]] = (x[:, col] >> (unit * k)) & 0xF + return res + + scales = scales.t().contiguous() + if format == "NN": + zeros = unpack_tensor(zeros, order_map=[0, 2, 4, 6, 1, 3, 5, 7]) + zeros = zeros.t().contiguous() + qweights = unpack_tensor(qweights, order_map=[0, 2, 4, 6, 1, 3, 5, 7]) + qweights = qweights.t().contiguous() + else: + zeros = unpack_tensor(zeros) + zeros = zeros.t().contiguous() + qweights = unpack_tensor(qweights) + output_dim, input_dim = qweights.shape + qweights = qweights.view(output_dim, input_dim // group_size, group_size) + zeros = zeros.view(output_dim, input_dim // group_size, 1) + scales = scales.view(output_dim, input_dim // group_size, 1) + + qweights = (qweights - zeros) * scales + qweights = qweights.view(output_dim, input_dim) + if only_return_weight: + return qweights + output = torch.nn.functional.linear(inputs, qweights.to(inputs.dtype)) + return output, qweights + + +def wui4a16_gemm( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + bias: "torch.Tensor" = None, + group_size: int = 128, + format: str = "NN", +): + output_shape = inputs.shape[:-1] + (scales.shape[1],) + + output = ops.infer.wui4a16_gemm( + inputs, qweights, scales, zeros, bias, group_size, format + ) + return output.view(output_shape) + + +def wui4a16_gemv( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + bias: "torch.Tensor" = None, + group_size: int = 128, + format: str = "NN", +): + output_shape = inputs.shape[:-1] + (scales.shape[1],) + + output = ops.infer.wui4a16_gemv( + inputs, qweights, scales, zeros, bias, group_size, format + ) + return output.view(output_shape) + + +def wui4a16( + inputs: "torch.Tensor", + qweights: "torch.Tensor", + scales: "torch.Tensor", + zeros: "torch.Tensor", + bias: "torch.Tensor" = None, + group_size: int = 128, + format: str = "NN", +): + """ + format = TN,NN + group_size = TN(128),NN(128, 32) + input : bfloat16|fp16 (bs, ic) + qweights : int32 NN: (ic, oc // 8) TN:(oc, ic // 8) + scales : bfloat16|fp16 (ic // group_size, oc) + zeros : int32 (ic // group_size, oc // 8) + bias : bfloat16|fp16 (oc, ) + output : bfloat16|fp16 (bs, oc) + 支持条件 : NN: oc % 8 == 0 && ic % group_size == 0 && ic % 2 == 0 + TN: oc % 2 == 0 && ic % group_size == 0 + """ + batch = inputs.numel() // inputs.shape[-1] + if batch <= 1: + return wui4a16_gemv( + inputs=inputs, + qweights=qweights, + scales=scales, + zeros=zeros, + bias=bias, + group_size=group_size, + format=format, + ) + else: + return wui4a16_gemm( + inputs=inputs, + qweights=qweights, + scales=scales, + zeros=zeros, + bias=bias, + group_size=group_size, + format=format, + ) diff --git a/ixformer_sdk/inference/models/__init__.py b/ixformer_sdk/inference/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/inference/models/clip/__init__.py b/ixformer_sdk/inference/models/clip/__init__.py new file mode 100644 index 00000000..5e4e4c65 --- /dev/null +++ b/ixformer_sdk/inference/models/clip/__init__.py @@ -0,0 +1 @@ +from .modeling_clip import CLIPModel diff --git a/ixformer_sdk/inference/models/clip/configuration_clip.py b/ixformer_sdk/inference/models/clip/configuration_clip.py new file mode 100644 index 00000000..1a406ae0 --- /dev/null +++ b/ixformer_sdk/inference/models/clip/configuration_clip.py @@ -0,0 +1,503 @@ +# coding=utf-8 +# Copyright 2021 The HuggingFace Inc. team. All rights reserved. +# +# 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. +""" CLIP model configuration""" + +import copy +import os +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Mapping, Optional, Union + +if TYPE_CHECKING: + from transformers.processing_utils import ProcessorMixin + from transformers.utils import TensorType + +from transformers.configuration_utils import PretrainedConfig +from transformers.onnx import OnnxConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = { + "openai/clip-vit-base-patch32": "https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/config.json", + # See all CLIP models at https://huggingface.co/models?filter=clip +} + + +class CLIPTextConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`CLIPTextModel`]. It is used to instantiate a CLIP + text encoder according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of the text encoder of the CLIP + [openai/clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + vocab_size (`int`, *optional*, defaults to 49408): + Vocabulary size of the CLIP text model. Defines the number of different tokens that can be represented by + the `inputs_ids` passed when calling [`CLIPModel`]. + hidden_size (`int`, *optional*, defaults to 512): + Dimensionality of the encoder layers and the pooler layer. + intermediate_size (`int`, *optional*, defaults to 2048): + Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. + num_hidden_layers (`int`, *optional*, defaults to 12): + Number of hidden layers in the Transformer encoder. + num_attention_heads (`int`, *optional*, defaults to 8): + Number of attention heads for each attention layer in the Transformer encoder. + max_position_embeddings (`int`, *optional*, defaults to 77): + The maximum sequence length that this model might ever be used with. Typically set this to something large + just in case (e.g., 512 or 1024 or 2048). + hidden_act (`str` or `function`, *optional*, defaults to `"quick_gelu"`): + The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, + `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported. + layer_norm_eps (`float`, *optional*, defaults to 1e-5): + The epsilon used by the layer normalization layers. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + initializer_factor (`float`, *optional*, defaults to 1): + A factor for initializing all weight matrices (should be kept to 1, used internally for initialization + testing). + + Example: + + ```python + >>> from transformers import CLIPTextConfig, CLIPTextModel + + >>> # Initializing a CLIPTextConfig with openai/clip-vit-base-patch32 style configuration + >>> configuration = CLIPTextConfig() + + >>> # Initializing a CLIPTextModel (with random weights) from the openai/clip-vit-base-patch32 style configuration + >>> model = CLIPTextModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + model_type = "clip_text_model" + + def __init__( + self, + vocab_size=49408, + hidden_size=512, + intermediate_size=2048, + projection_dim=512, + num_hidden_layers=12, + num_attention_heads=8, + max_position_embeddings=77, + hidden_act="quick_gelu", + layer_norm_eps=1e-5, + attention_dropout=0.0, + initializer_range=0.02, + initializer_factor=1.0, + pad_token_id=1, + bos_token_id=0, + eos_token_id=2, + **kwargs, + ): + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + **kwargs, + ) + + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.projection_dim = projection_dim + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.max_position_embeddings = max_position_embeddings + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.initializer_factor = initializer_factor + self.attention_dropout = attention_dropout + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs + ) -> "PretrainedConfig": + config_dict, kwargs = cls.get_config_dict( + pretrained_model_name_or_path, **kwargs + ) + + # get the text config dict if we are loading from CLIPConfig + if config_dict.get("model_type") == "clip": + config_dict = config_dict["text_config"] + + if ( + "model_type" in config_dict + and hasattr(cls, "model_type") + and config_dict["model_type"] != cls.model_type + ): + logger.warning( + f"You are using a model of type {config_dict['model_type']} to instantiate a model of type " + f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." + ) + + return cls.from_dict(config_dict, **kwargs) + + +class CLIPVisionConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`CLIPVisionModel`]. It is used to instantiate a + CLIP vision encoder according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the vision encoder of the CLIP + [openai/clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + hidden_size (`int`, *optional*, defaults to 768): + Dimensionality of the encoder layers and the pooler layer. + intermediate_size (`int`, *optional*, defaults to 3072): + Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. + num_hidden_layers (`int`, *optional*, defaults to 12): + Number of hidden layers in the Transformer encoder. + num_attention_heads (`int`, *optional*, defaults to 12): + Number of attention heads for each attention layer in the Transformer encoder. + image_size (`int`, *optional*, defaults to 224): + The size (resolution) of each image. + patch_size (`int`, *optional*, defaults to 32): + The size (resolution) of each patch. + hidden_act (`str` or `function`, *optional*, defaults to `"quick_gelu"`): + The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, + `"relu"`, `"selu"` and `"gelu_new"` ``"quick_gelu"` are supported. + layer_norm_eps (`float`, *optional*, defaults to 1e-5): + The epsilon used by the layer normalization layers. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + initializer_factor (`float`, *optional*, defaults to 1): + A factor for initializing all weight matrices (should be kept to 1, used internally for initialization + testing). + + Example: + + ```python + >>> from transformers import CLIPVisionConfig, CLIPVisionModel + + >>> # Initializing a CLIPVisionConfig with openai/clip-vit-base-patch32 style configuration + >>> configuration = CLIPVisionConfig() + + >>> # Initializing a CLIPVisionModel (with random weights) from the openai/clip-vit-base-patch32 style configuration + >>> model = CLIPVisionModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "clip_vision_model" + + def __init__( + self, + hidden_size=768, + intermediate_size=3072, + projection_dim=512, + num_hidden_layers=12, + num_attention_heads=12, + num_channels=3, + image_size=224, + patch_size=32, + hidden_act="quick_gelu", + layer_norm_eps=1e-5, + attention_dropout=0.0, + initializer_range=0.02, + initializer_factor=1.0, + **kwargs, + ): + super().__init__(**kwargs) + + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.projection_dim = projection_dim + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_channels = num_channels + self.patch_size = patch_size + self.image_size = image_size + self.initializer_range = initializer_range + self.initializer_factor = initializer_factor + self.attention_dropout = attention_dropout + self.layer_norm_eps = layer_norm_eps + self.hidden_act = hidden_act + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs + ) -> "PretrainedConfig": + config_dict, kwargs = cls.get_config_dict( + pretrained_model_name_or_path, **kwargs + ) + + # get the vision config dict if we are loading from CLIPConfig + if config_dict.get("model_type") == "clip": + config_dict = config_dict["vision_config"] + + if ( + "model_type" in config_dict + and hasattr(cls, "model_type") + and config_dict["model_type"] != cls.model_type + ): + logger.warning( + f"You are using a model of type {config_dict['model_type']} to instantiate a model of type " + f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." + ) + + return cls.from_dict(config_dict, **kwargs) + + +class CLIPConfig(PretrainedConfig): + r""" + [`CLIPConfig`] is the configuration class to store the configuration of a [`CLIPModel`]. It is used to instantiate + a CLIP model according to the specified arguments, defining the text model and vision model configs. Instantiating + a configuration with the defaults will yield a similar configuration to that of the CLIP + [openai/clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + text_config (`dict`, *optional*): + Dictionary of configuration options used to initialize [`CLIPTextConfig`]. + vision_config (`dict`, *optional*): + Dictionary of configuration options used to initialize [`CLIPVisionConfig`]. + projection_dim (`int`, *optional*, defaults to 512): + Dimentionality of text and vision projection layers. + logit_scale_init_value (`float`, *optional*, defaults to 2.6592): + The inital value of the *logit_scale* paramter. Default is used as per the original CLIP implementation. + kwargs (*optional*): + Dictionary of keyword arguments. + + Example: + + ```python + >>> from transformers import CLIPConfig, CLIPModel + + >>> # Initializing a CLIPConfig with openai/clip-vit-base-patch32 style configuration + >>> configuration = CLIPConfig() + + >>> # Initializing a CLIPModel (with random weights) from the openai/clip-vit-base-patch32 style configuration + >>> model = CLIPModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + + >>> # We can also initialize a CLIPConfig from a CLIPTextConfig and a CLIPVisionConfig + >>> from transformers import CLIPTextConfig, CLIPVisionConfig + + >>> # Initializing a CLIPText and CLIPVision configuration + >>> config_text = CLIPTextConfig() + >>> config_vision = CLIPVisionConfig() + + >>> config = CLIPConfig.from_text_vision_configs(config_text, config_vision) + ```""" + + model_type = "clip" + is_composition = True + + def __init__( + self, + text_config=None, + vision_config=None, + projection_dim=512, + logit_scale_init_value=2.6592, + **kwargs, + ): + # If `_config_dict` exist, we use them for the backward compatibility. + # We pop out these 2 attributes before calling `super().__init__` to avoid them being saved (which causes a lot + # of confusion!). + text_config_dict = kwargs.pop("text_config_dict", None) + vision_config_dict = kwargs.pop("vision_config_dict", None) + + super().__init__(**kwargs) + + # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in + # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most + # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`. + if text_config_dict is not None: + if text_config is None: + text_config = {} + + # This is the complete result when using `text_config_dict`. + _text_config_dict = CLIPTextConfig(**text_config_dict).to_dict() + + # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different. + for key, value in _text_config_dict.items(): + if ( + key in text_config + and value != text_config[key] + and key not in ["transformers_version"] + ): + # If specified in `text_config_dict` + if key in text_config_dict: + message = ( + f"`{key}` is found in both `text_config_dict` and `text_config` but with different values. " + f'The value `text_config_dict["{key}"]` will be used instead.' + ) + # If inferred from default argument values (just to be super careful) + else: + message = ( + f"`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The " + f'value `text_config["{key}"]` will be overriden.' + ) + logger.warning(message) + + # Update all values in `text_config` with the ones in `_text_config_dict`. + text_config.update(_text_config_dict) + + if vision_config_dict is not None: + if vision_config is None: + vision_config = {} + + # This is the complete result when using `vision_config_dict`. + _vision_config_dict = CLIPVisionConfig(**vision_config_dict).to_dict() + # convert keys to string instead of integer + if "id2label" in _vision_config_dict: + _vision_config_dict["id2label"] = { + str(key): value + for key, value in _vision_config_dict["id2label"].items() + } + + # Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different. + for key, value in _vision_config_dict.items(): + if ( + key in vision_config + and value != vision_config[key] + and key not in ["transformers_version"] + ): + # If specified in `vision_config_dict` + if key in vision_config_dict: + message = ( + f"`{key}` is found in both `vision_config_dict` and `vision_config` but with different " + f'values. The value `vision_config_dict["{key}"]` will be used instead.' + ) + # If inferred from default argument values (just to be super careful) + else: + message = ( + f"`vision_config_dict` is provided which will be used to initialize `CLIPVisionConfig`. " + f'The value `vision_config["{key}"]` will be overriden.' + ) + logger.warning(message) + + # Update all values in `vision_config` with the ones in `_vision_config_dict`. + vision_config.update(_vision_config_dict) + + if text_config is None: + text_config = {} + logger.info( + "`text_config` is `None`. Initializing the `CLIPTextConfig` with default values." + ) + + if vision_config is None: + vision_config = {} + logger.info( + "`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values." + ) + + self.text_config = CLIPTextConfig(**text_config) + self.vision_config = CLIPVisionConfig(**vision_config) + + self.projection_dim = projection_dim + self.logit_scale_init_value = logit_scale_init_value + self.initializer_factor = 1.0 + + @classmethod + def from_text_vision_configs( + cls, text_config: CLIPTextConfig, vision_config: CLIPVisionConfig, **kwargs + ): + r""" + Instantiate a [`CLIPConfig`] (or a derived class) from clip text model configuration and clip vision model + configuration. + + Returns: + [`CLIPConfig`]: An instance of a configuration object + """ + + return cls( + text_config=text_config.to_dict(), + vision_config=vision_config.to_dict(), + **kwargs, + ) + + def to_dict(self): + """ + Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`]. + + Returns: + `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance, + """ + output = copy.deepcopy(self.__dict__) + output["text_config"] = self.text_config.to_dict() + output["vision_config"] = self.vision_config.to_dict() + output["model_type"] = self.__class__.model_type + return output + + +class CLIPOnnxConfig(OnnxConfig): + @property + def inputs(self) -> Mapping[str, Mapping[int, str]]: + return OrderedDict( + [ + ("input_ids", {0: "batch", 1: "sequence"}), + ( + "pixel_values", + {0: "batch", 1: "num_channels", 2: "height", 3: "width"}, + ), + ("attention_mask", {0: "batch", 1: "sequence"}), + ] + ) + + @property + def outputs(self) -> Mapping[str, Mapping[int, str]]: + return OrderedDict( + [ + ("logits_per_image", {0: "batch"}), + ("logits_per_text", {0: "batch"}), + ("text_embeds", {0: "batch"}), + ("image_embeds", {0: "batch"}), + ] + ) + + @property + def atol_for_validation(self) -> float: + return 1e-4 + + def generate_dummy_inputs( + self, + processor: "ProcessorMixin", + batch_size: int = -1, + seq_length: int = -1, + framework: Optional["TensorType"] = None, + ) -> Mapping[str, Any]: + text_input_dict = super().generate_dummy_inputs( + processor.tokenizer, + batch_size=batch_size, + seq_length=seq_length, + framework=framework, + ) + image_input_dict = super().generate_dummy_inputs( + processor.feature_extractor, batch_size=batch_size, framework=framework + ) + return {**text_input_dict, **image_input_dict} + + @property + def default_onnx_opset(self) -> int: + return 14 diff --git a/ixformer_sdk/inference/models/clip/modeling_clip.py b/ixformer_sdk/inference/models/clip/modeling_clip.py new file mode 100644 index 00000000..cb98ea23 --- /dev/null +++ b/ixformer_sdk/inference/models/clip/modeling_clip.py @@ -0,0 +1,1578 @@ +# coding=utf-8 +# Copyright 2021 The OpenAI Team Authors and The HuggingFace Team. All rights reserved. +# +# 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. +""" PyTorch CLIP model.""" +import ixformer.functions as ixf_F + +using_ixf_linear = True # 提升1.3-1.5倍 +using_ixf_bmm = False # 无效 +using_ixf_layernorm = False # 无效 +using_ixf_conv2d = False # + + +from dataclasses import dataclass +from typing import Any, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn +from transformers.activations import ACT2FN +from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ( + ModelOutput, + add_start_docstrings, + add_start_docstrings_to_model_forward, + logging, + replace_return_docstrings, +) + +from .configuration_clip import CLIPConfig, CLIPTextConfig, CLIPVisionConfig + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "openai/clip-vit-base-patch32" + +CLIP_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "openai/clip-vit-base-patch32", + # See all CLIP models at https://huggingface.co/models?filter=clip +] + + +# Copied from transformers.models.bart.modeling_bart._expand_mask +def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(dtype).min + ) + + +# contrastive loss function, adapted from +# https://sachinruk.github.io/blog/pytorch/pytorch%20lightning/loss%20function/gpu/2021/03/07/CLIP.html +def contrastive_loss(logits: torch.Tensor) -> torch.Tensor: + return nn.functional.cross_entropy( + logits, torch.arange(len(logits), device=logits.device) + ) + + +def clip_loss(similarity: torch.Tensor) -> torch.Tensor: + caption_loss = contrastive_loss(similarity) + image_loss = contrastive_loss(similarity.t()) + return (caption_loss + image_loss) / 2.0 + + +@dataclass +class CLIPVisionModelOutput(ModelOutput): + """ + Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states. + + Args: + image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The image embeddings obtained by applying the projection layer to the pooler_output. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + image_embeds: Optional[torch.FloatTensor] = None + last_hidden_state: torch.FloatTensor = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class CLIPTextModelOutput(ModelOutput): + """ + Base class for text model's outputs that also contains a pooling of the last hidden states. + + Args: + text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`): + The text embeddings obtained by applying the projection layer to the pooler_output. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + text_embeds: Optional[torch.FloatTensor] = None + last_hidden_state: torch.FloatTensor = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class CLIPOutput(ModelOutput): + """ + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): + Contrastive loss for image-text similarity. + logits_per_image:(`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): + The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text + similarity scores. + logits_per_text:(`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): + The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image + similarity scores. + text_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`): + The text embeddings obtained by applying the projection layer to the pooled output of [`CLIPTextModel`]. + image_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`): + The image embeddings obtained by applying the projection layer to the pooled output of [`CLIPVisionModel`]. + text_model_output(`BaseModelOutputWithPooling`): + The output of the [`CLIPTextModel`]. + vision_model_output(`BaseModelOutputWithPooling`): + The output of the [`CLIPVisionModel`]. + """ + + loss: Optional[torch.FloatTensor] = None + logits_per_image: torch.FloatTensor = None + logits_per_text: torch.FloatTensor = None + text_embeds: torch.FloatTensor = None + image_embeds: torch.FloatTensor = None + text_model_output: BaseModelOutputWithPooling = None + vision_model_output: BaseModelOutputWithPooling = None + + def to_tuple(self) -> Tuple[Any]: + return tuple( + self[k] + if k not in ["text_model_output", "vision_model_output"] + else getattr(self, k).to_tuple() + for k in self.keys() + ) + + +class CLIPVisionEmbeddings(nn.Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + self.using_nhwc = False + if self.using_nhwc: + self.patch_embedding = self.patch_embedding.to( + memory_format=torch.channels_last + ) + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer( + "position_ids", torch.arange(self.num_positions).expand((1, -1)) + ) + + def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor: + batch_size = pixel_values.shape[0] + + if using_ixf_conv2d: + pixel_values = pixel_values.permute(0, 2, 3, 1).contiguous() + weights = self.patch_embedding.weight.permute(0, 2, 3, 1).contiguous() + patch_embeds = ixf_F.conv2d( + pixel_values, weights, None, stride=self.patch_size + ) + patch_embeds = patch_embeds.permute(0, 3, 1, 2).contiguous() + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + elif self.using_nhwc: # pytorch conv2d nhwc + pixel_values = pixel_values.to(memory_format=torch.channels_last) + + patch_embeds = self.patch_embedding(pixel_values) + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + + else: + patch_embeds = self.patch_embedding( + pixel_values + ) # shape = [*, width, grid, grid] + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + + class_embeds = self.class_embedding.expand(batch_size, 1, -1) + embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + embeddings = embeddings + self.position_embedding(self.position_ids) + return embeddings + + +class CLIPTextEmbeddings(nn.Module): + def __init__(self, config: CLIPTextConfig): + super().__init__() + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding( + config.max_position_embeddings, embed_dim + ) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)) + ) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.Tensor: + seq_length = ( + input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + ) + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + +class CLIPAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return ( + tensor.view(bsz, seq_len, self.num_heads, self.head_dim) + .transpose(1, 2) + .contiguous() + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + causal_attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + """Input shape: Batch x Time x Channel""" + + bsz, tgt_len, embed_dim = hidden_states.size() + using_ixf_linear = True + # get query proj + if using_ixf_linear: + query_states = ( + ixf_F.linear(hidden_states, self.q_proj.weight, self.q_proj.bias) + * self.scale + ) + else: + query_states = self.q_proj(hidden_states) * self.scale + if using_ixf_linear: + key_states = self._shape( + ixf_F.linear(hidden_states, self.k_proj.weight, self.k_proj.bias), + -1, + bsz, + ) + else: + key_states = self._shape(self.k_proj(hidden_states), -1, bsz) + if using_ixf_linear: + value_states = self._shape( + ixf_F.linear(hidden_states, self.v_proj.weight, self.v_proj.bias), + -1, + bsz, + ) + else: + value_states = self._shape(self.v_proj(hidden_states), -1, bsz) + + proj_shape = (bsz * self.num_heads, -1, self.head_dim) + query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) + key_states = key_states.view(*proj_shape) + value_states = value_states.view(*proj_shape) + + src_len = key_states.size(1) + + if using_ixf_bmm: + attn_weights = ixf_F.act_bias_mm( + query_states, key_states, scale=1, trans_format="TN" + ) + else: + attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) + if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): + raise ValueError( + f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is" + f" {attn_weights.size()}" + ) + + # apply the causal_attention_mask first + if causal_attention_mask is not None: + if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is" + f" {causal_attention_mask.size()}" + ) + attn_weights = ( + attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + + causal_attention_mask + ) + attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, tgt_len, src_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}" + ) + attn_weights = ( + attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + + attention_mask + ) + attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + if output_attentions: + # this operation is a bit akward, but it's required to + # make sure that attn_weights keeps its gradient. + # In order to do so, attn_weights have to reshaped + # twice and have to be reused in the following + attn_weights_reshaped = attn_weights.view( + bsz, self.num_heads, tgt_len, src_len + ) + attn_weights = attn_weights_reshaped.view( + bsz * self.num_heads, tgt_len, src_len + ) + else: + attn_weights_reshaped = None + + attn_probs = nn.functional.dropout( + attn_weights, p=self.dropout, training=self.training + ) + if using_ixf_bmm: + attn_output = ixf_F.act_bias_mm(attn_probs, value_states, trans_format="NN") + else: + attn_output = torch.bmm(attn_probs, value_states) + + if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim) + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, tgt_len, embed_dim) + # using_ixf_linear=False + if using_ixf_linear: + attn_output = ixf_F.linear( + attn_output, self.out_proj.weight, self.out_proj.bias + ) + else: + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights_reshaped + + +class CLIPMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if using_ixf_linear: + # hidden_states = ixf_F.linear(hidden_states, self.fc1.weight, self.fc1.bias) + input_shape = list(hidden_states.shape) + hidden_states = hidden_states.view(-1, input_shape[-1]) + hidden_states = ixf_F.act_bias_mm( + hidden_states, + self.fc1.weight, + self.fc1.bias, + scale=1, + act_type="gelu", + trans_format="TN", + ) + input_shape[-1] = -1 + hidden_states = hidden_states.view(*input_shape) + else: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + + if using_ixf_linear: + hidden_states = ixf_F.linear(hidden_states, self.fc2.weight, self.fc2.bias) + else: + hidden_states = self.fc2(hidden_states) + + return hidden_states + + +class CLIPEncoderLayer(nn.Module): + def __init__(self, config: CLIPConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = CLIPAttention(config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = CLIPMLP(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + causal_attention_mask: torch.Tensor, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.FloatTensor]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + `(config.encoder_attention_heads,)`. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + + if using_ixf_layernorm: + hidden_states = ixf_F.layernorm( + hidden_states, self.layer_norm1.weight, self.layer_norm1.bias + ) + else: + hidden_states = self.layer_norm1(hidden_states) + hidden_states, attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + causal_attention_mask=causal_attention_mask, + output_attentions=output_attentions, + ) + if using_ixf_layernorm: + hidden_states = ixf_F.residual_bias(hidden_states, residual) + else: + hidden_states = residual + hidden_states + + residual = hidden_states + if using_ixf_layernorm: + hidden_states = ixf_F.layernorm( + hidden_states, self.layer_norm2.weight, self.layer_norm2.bias + ) + else: + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + if using_ixf_layernorm: + hidden_states = ixf_F.residual_bias(hidden_states, residual) + else: + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + +class CLIPPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = CLIPConfig + base_model_prefix = "clip" + supports_gradient_checkpointing = True + _keys_to_ignore_on_load_missing = [r"position_ids"] + + def _init_weights(self, module): + """Initialize the weights""" + factor = self.config.initializer_factor + if isinstance(module, CLIPTextEmbeddings): + module.token_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02) + module.position_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02) + elif isinstance(module, CLIPVisionEmbeddings): + factor = self.config.initializer_factor + nn.init.normal_( + module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor + ) + nn.init.normal_( + module.patch_embedding.weight, + std=module.config.initializer_range * factor, + ) + nn.init.normal_( + module.position_embedding.weight, + std=module.config.initializer_range * factor, + ) + elif isinstance(module, CLIPAttention): + factor = self.config.initializer_factor + in_proj_std = ( + (module.embed_dim**-0.5) + * ((2 * module.config.num_hidden_layers) ** -0.5) + * factor + ) + out_proj_std = (module.embed_dim**-0.5) * factor + nn.init.normal_(module.q_proj.weight, std=in_proj_std) + nn.init.normal_(module.k_proj.weight, std=in_proj_std) + nn.init.normal_(module.v_proj.weight, std=in_proj_std) + nn.init.normal_(module.out_proj.weight, std=out_proj_std) + elif isinstance(module, CLIPMLP): + factor = self.config.initializer_factor + in_proj_std = ( + (module.config.hidden_size**-0.5) + * ((2 * module.config.num_hidden_layers) ** -0.5) + * factor + ) + fc_std = (2 * module.config.hidden_size) ** -0.5 * factor + nn.init.normal_(module.fc1.weight, std=fc_std) + nn.init.normal_(module.fc2.weight, std=in_proj_std) + elif isinstance(module, CLIPModel): + nn.init.normal_( + module.text_projection.weight, + std=module.text_embed_dim**-0.5 * self.config.initializer_factor, + ) + nn.init.normal_( + module.visual_projection.weight, + std=module.vision_embed_dim**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, CLIPVisionModelWithProjection): + nn.init.normal_( + module.visual_projection.weight, + std=self.config.hidden_size**-0.5 * self.config.initializer_factor, + ) + elif isinstance(module, CLIPTextModelWithProjection): + nn.init.normal_( + module.text_projection.weight, + std=self.config.hidden_size**-0.5 * self.config.initializer_factor, + ) + + if isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + if isinstance(module, nn.Linear) and module.bias is not None: + module.bias.data.zero_() + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, CLIPEncoder): + module.gradient_checkpointing = value + + +CLIP_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`CLIPConfig`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +CLIP_TEXT_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + +CLIP_VISION_INPUTS_DOCSTRING = r""" + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using + [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + +CLIP_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using + [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details. + return_loss (`bool`, *optional*): + Whether or not to return the contrastive loss. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +class CLIPEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`CLIPEncoderLayer`]. + + Args: + config: CLIPConfig + """ + + def __init__(self, config: CLIPConfig): + super().__init__() + self.config = config + self.layers = nn.ModuleList( + [CLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.gradient_checkpointing = False + + def forward( + self, + inputs_embeds, + attention_mask: Optional[torch.Tensor] = None, + causal_attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. + This is useful if you want more control over how to convert `input_ids` indices into associated vectors + than the model's internal embedding lookup matrix. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Causal mask for the text model. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + hidden_states = inputs_embeds + for idx, encoder_layer in enumerate(self.layers): + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs, output_attentions) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(encoder_layer), + hidden_states, + attention_mask, + causal_attention_mask, + ) + else: + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + causal_attention_mask, + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, encoder_states, all_attentions] + if v is not None + ) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=encoder_states, + attentions=all_attentions, + ) + + +class CLIPTextTransformer(nn.Module): + def __init__(self, config: CLIPTextConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + self.embeddings = CLIPTextEmbeddings(config) + self.encoder = CLIPEncoder(config) + self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=BaseModelOutputWithPooling, config_class=CLIPTextConfig + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + """ + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if input_ids is None: + raise ValueError("You have to specify input_ids") + + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + + hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) + + bsz, seq_len = input_shape + # CLIP's text model uses causal mask, prepare it here. + # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 + causal_attention_mask = self._build_causal_attention_mask( + bsz, seq_len, hidden_states.dtype + ).to(hidden_states.device) + # expand attention_mask + if attention_mask is not None: + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + attention_mask = _expand_mask(attention_mask, hidden_states.dtype) + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + attention_mask=attention_mask, + causal_attention_mask=causal_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[0] + last_hidden_state = self.final_layer_norm(last_hidden_state) + + # text_embeds.shape = [batch_size, sequence_length, transformer.width] + # take features from the eot embedding (eot_token is the highest number in each sequence) + # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax( + dim=-1 + ), + ] + + if not return_dict: + return (last_hidden_state, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + def _build_causal_attention_mask(self, bsz, seq_len, dtype): + # lazily create causal attention mask, with full attention between the vision tokens + # pytorch uses additive attention mask; fill with -inf + mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype) + mask.fill_(torch.tensor(torch.finfo(dtype).min)) + mask.triu_(1) # zero out the lower diagonal + mask = mask.unsqueeze(1) # expand mask + return mask + + +@add_start_docstrings( + """The text model from CLIP without any head or projection on top.""", + CLIP_START_DOCSTRING, +) +class CLIPTextModel(CLIPPreTrainedModel): + config_class = CLIPTextConfig + + _no_split_modules = ["CLIPEncoderLayer"] + + def __init__(self, config: CLIPTextConfig): + super().__init__(config) + self.text_model = CLIPTextTransformer(config) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.text_model.embeddings.token_embedding + + def set_input_embeddings(self, value): + self.text_model.embeddings.token_embedding = value + + @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=BaseModelOutputWithPooling, config_class=CLIPTextConfig + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + Examples: + + ```python + >>> from transformers import AutoTokenizer, CLIPTextModel + + >>> model = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32") + >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled (EOS token) states + ```""" + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + return self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + +class CLIPVisionTransformer(nn.Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPVisionEmbeddings(config) + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.encoder = CLIPEncoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=BaseModelOutputWithPooling, config_class=CLIPVisionConfig + ) + def forward( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + """ + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + hidden_states = self.embeddings(pixel_values) + hidden_states = self.pre_layrnorm(hidden_states) + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[0] + pooled_output = last_hidden_state[:, 0, :] + pooled_output = self.post_layernorm(pooled_output) + + if not return_dict: + return (last_hidden_state, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + +@add_start_docstrings( + """The vision model from CLIP without any head or projection on top.""", + CLIP_START_DOCSTRING, +) +class CLIPVisionModel(CLIPPreTrainedModel): + config_class = CLIPVisionConfig + main_input_name = "pixel_values" + + def __init__(self, config: CLIPVisionConfig): + super().__init__(config) + self.vision_model = CLIPVisionTransformer(config) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.vision_model.embeddings.patch_embedding + + @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=BaseModelOutputWithPooling, config_class=CLIPVisionConfig + ) + def forward( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, CLIPVisionModel + + >>> model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> last_hidden_state = outputs.last_hidden_state + >>> pooled_output = outputs.pooler_output # pooled CLS states + ```""" + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + return self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + +@add_start_docstrings(CLIP_START_DOCSTRING) +class CLIPModel(CLIPPreTrainedModel): + config_class = CLIPConfig + + def __init__(self, config: CLIPConfig): + super().__init__(config) + + if not isinstance(config.text_config, CLIPTextConfig): + raise ValueError( + "config.text_config is expected to be of type CLIPTextConfig but is of type" + f" {type(config.text_config)}." + ) + + if not isinstance(config.vision_config, CLIPVisionConfig): + raise ValueError( + "config.vision_config is expected to be of type CLIPVisionConfig but is of type" + f" {type(config.vision_config)}." + ) + + text_config = config.text_config + vision_config = config.vision_config + + self.projection_dim = config.projection_dim + self.text_embed_dim = text_config.hidden_size + self.vision_embed_dim = vision_config.hidden_size + + self.text_model = CLIPTextTransformer(text_config) + self.vision_model = CLIPVisionTransformer(vision_config) + + self.visual_projection = nn.Linear( + self.vision_embed_dim, self.projection_dim, bias=False + ) + self.text_projection = nn.Linear( + self.text_embed_dim, self.projection_dim, bias=False + ) + self.logit_scale = nn.Parameter( + torch.ones([]) * self.config.logit_scale_init_value + ) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING) + def get_text_features( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> torch.FloatTensor: + r""" + Returns: + text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by + applying the projection layer to the pooled output of [`CLIPTextModel`]. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, CLIPModel + + >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + >>> text_features = model.get_text_features(**inputs) + ```""" + # Use CLIP model's config for some fields (if specified) instead of those of vision & text components. + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + text_outputs = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = text_outputs[1] + text_features = self.text_projection(pooled_output) + + return text_features + + @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING) + def get_image_features( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> torch.FloatTensor: + r""" + Returns: + image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by + applying the projection layer to the pooled output of [`CLIPVisionModel`]. + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, CLIPModel + + >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> image_features = model.get_image_features(**inputs) + ```""" + # Use CLIP model's config for some fields (if specified) instead of those of vision & text components. + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + vision_outputs = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = vision_outputs[1] # pooled_output + if using_ixf_linear: + image_features = ixf_F.linear( + pooled_output, + self.visual_projection.weight, + self.visual_projection.bias, + ) + else: + image_features = self.visual_projection(pooled_output) + + return image_features + + @add_start_docstrings_to_model_forward(CLIP_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CLIPOutput, config_class=CLIPConfig) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + return_loss: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CLIPOutput]: + r""" + Returns: + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, CLIPModel + + >>> model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor( + ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True + ... ) + + >>> outputs = model(**inputs) + >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score + >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities + ```""" + # Use CLIP model's config for some fields (if specified) instead of those of vision & text components. + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + vision_outputs = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + text_outputs = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + image_embeds = vision_outputs[1] + if using_ixf_linear: + image_embeds = ixf_F.linear( + image_embeds, self.visual_projection.weight, self.visual_projection.bias + ) + else: + image_embeds = self.visual_projection(image_embeds) + + text_embeds = text_outputs[1] + if using_ixf_linear: + text_embeds = ixf_F.linear( + text_embeds, self.text_projection.weight, self.text_projection.bias + ) + else: + text_embeds = self.text_projection(text_embeds) + + # normalized features + image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True) + text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True) + + # cosine similarity as logits + logit_scale = self.logit_scale.exp() + logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale + logits_per_image = logits_per_text.t() + + loss = None + if return_loss: + loss = clip_loss(logits_per_text) + + if not return_dict: + output = ( + logits_per_image, + logits_per_text, + text_embeds, + image_embeds, + text_outputs, + vision_outputs, + ) + return ((loss,) + output) if loss is not None else output + + return CLIPOutput( + loss=loss, + logits_per_image=logits_per_image, + logits_per_text=logits_per_text, + text_embeds=text_embeds, + image_embeds=image_embeds, + text_model_output=text_outputs, + vision_model_output=vision_outputs, + ) + + +@add_start_docstrings( + """ + CLIP Text Model with a projection layer on top (a linear layer on top of the pooled output). + """, + CLIP_START_DOCSTRING, +) +class CLIPTextModelWithProjection(CLIPPreTrainedModel): + config_class = CLIPTextConfig + + _no_split_modules = ["CLIPEncoderLayer"] + + def __init__(self, config: CLIPTextConfig): + super().__init__(config) + + self.text_model = CLIPTextTransformer(config) + + self.text_projection = nn.Linear( + config.hidden_size, config.projection_dim, bias=False + ) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.text_model.embeddings.token_embedding + + def set_input_embeddings(self, value): + self.text_model.embeddings.token_embedding = value + + @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=CLIPTextModelOutput, config_class=CLIPTextConfig + ) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CLIPTextModelOutput]: + r""" + Returns: + + Examples: + + ```python + >>> from transformers import AutoTokenizer, CLIPTextModelWithProjection + + >>> model = CLIPTextModelWithProjection.from_pretrained("openai/clip-vit-base-patch32") + >>> tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32") + + >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> text_embeds = outputs.text_embeds + ```""" + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + text_outputs = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = text_outputs[1] + if using_ixf_linear: + text_embeds = ixf_F.linear( + pooled_output, self.text_projection.weight, self.text_projection.bias + ) + else: + text_embeds = self.text_projection(pooled_output) + + if not return_dict: + outputs = (text_embeds, text_outputs[0]) + text_outputs[2:] + return tuple(output for output in outputs if output is not None) + + return CLIPTextModelOutput( + text_embeds=text_embeds, + last_hidden_state=text_outputs.last_hidden_state, + hidden_states=text_outputs.hidden_states, + attentions=text_outputs.attentions, + ) + + +@add_start_docstrings( + """ + CLIP Vision Model with a projection layer on top (a linear layer on top of the pooled output). + """, + CLIP_START_DOCSTRING, +) +class CLIPVisionModelWithProjection(CLIPPreTrainedModel): + config_class = CLIPVisionConfig + main_input_name = "pixel_values" + + def __init__(self, config: CLIPVisionConfig): + super().__init__(config) + + self.vision_model = CLIPVisionTransformer(config) + + self.visual_projection = nn.Linear( + config.hidden_size, config.projection_dim, bias=False + ) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> nn.Module: + return self.vision_model.embeddings.patch_embedding + + @add_start_docstrings_to_model_forward(CLIP_VISION_INPUTS_DOCSTRING) + @replace_return_docstrings( + output_type=CLIPVisionModelOutput, config_class=CLIPVisionConfig + ) + def forward( + self, + pixel_values: Optional[torch.FloatTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CLIPVisionModelOutput]: + r""" + Returns: + + Examples: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, CLIPVisionModelWithProjection + + >>> model = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-base-patch32") + >>> processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32") + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> image_embeds = outputs.image_embeds + ```""" + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + vision_outputs = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + pooled_output = vision_outputs[1] # pooled_output + + image_embeds = self.visual_projection(pooled_output) + + if not return_dict: + outputs = (image_embeds, vision_outputs[0]) + vision_outputs[2:] + return tuple(output for output in outputs if output is not None) + + return CLIPVisionModelOutput( + image_embeds=image_embeds, + last_hidden_state=vision_outputs.last_hidden_state, + hidden_states=vision_outputs.hidden_states, + attentions=vision_outputs.attentions, + ) diff --git a/ixformer_sdk/inference/models/codeshell_7b_chat/__init__.py b/ixformer_sdk/inference/models/codeshell_7b_chat/__init__.py new file mode 100644 index 00000000..a47c9774 --- /dev/null +++ b/ixformer_sdk/inference/models/codeshell_7b_chat/__init__.py @@ -0,0 +1 @@ +from .modeling_codeshell import CodeShellForCausalLM diff --git a/ixformer_sdk/inference/models/codeshell_7b_chat/configuration_codeshell.py b/ixformer_sdk/inference/models/codeshell_7b_chat/configuration_codeshell.py new file mode 100644 index 00000000..b95bca7b --- /dev/null +++ b/ixformer_sdk/inference/models/codeshell_7b_chat/configuration_codeshell.py @@ -0,0 +1,153 @@ +# coding=utf-8 +# Copyright 2023 WisdomShell Inc. All Rights Reserved. + +# 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. + +# This code is based on Bigcode's GPTBigCode configuration. It has been modified from +# its original forms to accommodate minor architectural differences compared to +# GPTBigCode Configuration that trained the model. + +# coding=utf-8 +# Copyright 2023 The BigCode team and HuggingFace Inc. 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. +""" CodeShell configuration""" + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +class CodeShellConfig(PretrainedConfig): + """ + This is the configuration class to store the configuration of a [`CodeShellModel`]. It is used to instantiate a + CodeShell model according to the specified arguments, defining the model architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + vocab_size (`int`, *optional*, defaults to 50257): + Vocabulary size of the GPT-2 model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`CodeShellModel`]. + n_positions (`int`, *optional*, defaults to 1024): + The maximum sequence length that this model might ever be used with. Typically set this to something large + just in case (e.g., 512 or 1024 or 2048). + n_embd (`int`, *optional*, defaults to 768): + Dimensionality of the embeddings and hidden states. + n_layer (`int`, *optional*, defaults to 12): + Number of hidden layers in the Transformer encoder. + n_head (`int`, *optional*, defaults to 12): + Number of attention heads for each attention layer in the Transformer encoder. + n_inner (`int`, *optional*, defaults to None): + Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd + activation_function (`str`, *optional*, defaults to `"gelu_pytorch_tanh"`): + Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new", + "gelu_pytorch_tanh"]`. + resid_pdrop (`float`, *optional*, defaults to 0.1): + The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. + embd_pdrop (`float`, *optional*, defaults to 0.1): + The dropout ratio for the embeddings. + attn_pdrop (`float`, *optional*, defaults to 0.1): + The dropout ratio for the attention. + layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): + The epsilon to use in the layer normalization layers. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + scale_attn_weights (`bool`, *optional*, defaults to `True`): + Scale attention weights by dividing by sqrt(hidden_size).. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). + attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`): + Whether to call the fused softmax in float32. + scale_attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`): + Whether to scale the attention softmax in float32. + attention_type (`bool`, *optional*, defaults to `True`): + Whether to use Multi-Query Attion (`True`) or Multi-Head Attention (`False`). + """ + + model_type = "codeshell" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "hidden_size": "n_embd", + "max_position_embeddings": "n_positions", + "num_attention_heads": "n_head", + "num_hidden_layers": "n_layer", + } + + def __init__( + self, + vocab_size=70144, + n_positions=8192, + n_embd=4096, + n_layer=42, + n_head=32, + n_inner=None, + activation_function="gelu_pytorch_tanh", + resid_pdrop=0.1, + embd_pdrop=0.1, + attn_pdrop=0.1, + layer_norm_epsilon=1e-5, + initializer_range=0.02, + scale_attn_weights=True, + use_cache=True, + bos_token_id=70000, + eos_token_id=70000, + attention_softmax_in_fp32=True, + scale_attention_softmax_in_fp32=True, + group_query_attention=True, + num_query_groups=1, + position_embedding_type="learned_absolute", + rope_scaling=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.n_positions = n_positions + self.n_embd = n_embd + self.n_layer = n_layer + self.n_head = n_head + self.n_inner = n_inner + self.activation_function = activation_function + self.resid_pdrop = resid_pdrop + self.embd_pdrop = embd_pdrop + self.attn_pdrop = attn_pdrop + self.layer_norm_epsilon = layer_norm_epsilon + self.initializer_range = initializer_range + self.scale_attn_weights = scale_attn_weights + self.use_cache = use_cache + self.attention_softmax_in_fp32 = attention_softmax_in_fp32 + self.scale_attention_softmax_in_fp32 = scale_attention_softmax_in_fp32 + self.group_query_attention = group_query_attention + self.num_query_groups = num_query_groups + self.position_embedding_type = position_embedding_type + self.rope_scaling = rope_scaling + assert self.position_embedding_type in [ + "learned_absolute", + "rope", + ], "position_embedding_type must be one of ['learned_absolute', 'rope']" + + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + + super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) diff --git a/ixformer_sdk/inference/models/codeshell_7b_chat/modeling_codeshell.py b/ixformer_sdk/inference/models/codeshell_7b_chat/modeling_codeshell.py new file mode 100644 index 00000000..50695a20 --- /dev/null +++ b/ixformer_sdk/inference/models/codeshell_7b_chat/modeling_codeshell.py @@ -0,0 +1,1280 @@ +# coding=utf-8 +# Copyright 2023 WisdomShell Inc. All Rights Reserved. + +# 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. + +# This code is based on Bigcode's GPTBigCode model. It has been modified from +# its original forms to accommodate minor architectural differences compared to +# GPTBigCode model that trained the model. + +# Copyright 2023 The Bigcode team and HuggingFace Inc. 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. +"""PyTorch CodeShell model.""" +import math +import os +from queue import Queue +from threading import Thread +from typing import Callable, List, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from transformers import ( + LogitsProcessorList, + PretrainedConfig, + PreTrainedModel, + StoppingCriteria, + StoppingCriteriaList, +) +from transformers.activations import ACT2FN +from transformers.generation.utils import GenerationConfig +from transformers.modeling_outputs import ( + BaseModelOutputWithPastAndCrossAttentions, + CausalLMOutputWithCrossAttentions, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, +) + +from .configuration_codeshell import CodeShellConfig +from .modeling_codeshell_ixformer import mha, mlp_forward + + +# Fused kernels +# Use separate functions for each case because conditionals prevent kernel fusion. +# TODO: Could have better fused kernels depending on scaling, dropout and head mask. +# Is it doable without writing 32 functions? +@torch.jit.script +def upcast_masked_softmax( + x: torch.Tensor, + mask: torch.Tensor, + mask_value: torch.Tensor, + scale: float, + softmax_dtype: torch.dtype, +): + input_dtype = x.dtype + x = x.to(softmax_dtype) * scale + x = torch.where(mask, x, mask_value) + x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype) + return x + + +@torch.jit.script +def upcast_softmax(x: torch.Tensor, scale: float, softmax_dtype: torch.dtype): + input_dtype = x.dtype + x = x.to(softmax_dtype) * scale + x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype) + return x + + +@torch.jit.script +def masked_softmax(x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor): + x = torch.where(mask, x, mask_value) + x = torch.nn.functional.softmax(x, dim=-1) + return x + + +class CodeShellRotaryEmbedding(torch.nn.Module): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + super().__init__() + + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / ( + self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim) + ) + self.register_buffer("inv_freq", inv_freq) + + # Build here to make `torch.jit.trace` work. + self._set_cos_sin_cache( + seq_len=max_position_embeddings, + device=self.inv_freq.device, + dtype=torch.get_default_dtype(), + ) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange( + self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype + ) + + freqs = torch.einsum("i,j->ij", t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer( + "cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False + ) + self.register_buffer( + "sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False + ) + + def forward(self, x, seq_len=None): + # x: [bs, num_attention_heads, seq_len, head_size] + if seq_len > self.max_seq_len_cached: + self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) + + return ( + self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype), + self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype), + ) + + +class CodeShellLinearScalingRotaryEmbedding(CodeShellRotaryEmbedding): + """CodeShellRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev""" + + def __init__( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + ): + self.scaling_factor = scaling_factor + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange( + self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype + ) + t = t / self.scaling_factor + + freqs = torch.einsum("i,j->ij", t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer( + "cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False + ) + self.register_buffer( + "sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False + ) + + +class CodeShellDynamicNTKScalingRotaryEmbedding(CodeShellRotaryEmbedding): + """ShellRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla""" + + def __init__( + self, + dim, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + ): + self.scaling_factor = scaling_factor + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + + if seq_len > self.max_position_embeddings: + base = self.base * ( + (self.scaling_factor * seq_len / self.max_position_embeddings) + - (self.scaling_factor - 1) + ) ** (self.dim / (self.dim - 2)) + inv_freq = 1.0 / ( + base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim) + ) + self.register_buffer("inv_freq", inv_freq) + + t = torch.arange( + self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype + ) + + freqs = torch.einsum("i,j->ij", t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer( + "cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False + ) + self.register_buffer( + "sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False + ) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids): + # The first two dimensions of cos and sin are always 1, so we can `squeeze` them. + cos = cos.squeeze(1).squeeze(0) # [seq_len, dim] + sin = sin.squeeze(1).squeeze(0) # [seq_len, dim] + cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] + sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class CodeShellAttention(nn.Module): + def __init__(self, config, layer_idx=None): + super().__init__() + self.mask_value = None + + self.position_embedding_type = config.position_embedding_type + self.rope_scaling = config.rope_scaling + self.max_position_embeddings = config.max_position_embeddings + + self.group_query_attention = config.group_query_attention + self.num_query_groups = config.num_query_groups + self.num_key_value_groups = ( + config.num_attention_heads // config.num_query_groups + ) + + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + self.kv_heads = ( + config.num_query_groups if self.group_query_attention else self.num_heads + ) + self.kv_dim = self.kv_heads * self.head_dim + self.split_size = self.embed_dim + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + + self.layer_idx = layer_idx + + self.c_attn = nn.Linear(self.embed_dim, self.embed_dim + 2 * self.kv_dim) + self.c_proj = nn.Linear(self.embed_dim, self.embed_dim) + + self.attn_dropout = nn.Dropout(config.attn_pdrop) + self.resid_dropout = nn.Dropout(config.resid_pdrop) + + if self.position_embedding_type == "rope": + self._init_rope() + + def _init_rope(self): + if self.rope_scaling is None: + self.rotary_emb = CodeShellRotaryEmbedding( + self.head_dim, max_position_embeddings=self.max_position_embeddings + ) + else: + scaling_type = self.rope_scaling["type"] + scaling_factor = self.rope_scaling["factor"] + if scaling_type == "linear": + self.rotary_emb = CodeShellLinearScalingRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + ) + elif scaling_type == "dynamic": + self.rotary_emb = CodeShellDynamicNTKScalingRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + ) + else: + raise ValueError(f"Unknown RoPE scaling type {scaling_type}") + + def _get_mask_value(self, device, dtype): + # torch.where expects a tensor. We use a cache to avoid recreating it every time. + if ( + self.mask_value is None + or self.mask_value.dtype != dtype + or self.mask_value.device != device + ): + self.mask_value = torch.full( + [], torch.finfo(dtype).min, dtype=dtype, device=device + ) + return self.mask_value + + def forward( + self, + hidden_states: torch.Tensor, + layer_past: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + ) -> Union[ + Tuple[torch.Tensor, Optional[torch.Tensor]], + Tuple[torch.Tensor, Optional[torch.Tensor], Tuple[torch.Tensor, ...]], + ]: + bsz, q_len, _ = hidden_states.size() + query_states, key_states, value_states = self.c_attn(hidden_states).split( + (self.embed_dim, self.kv_dim, self.kv_dim), dim=2 + ) + + query_states = query_states.view( + bsz, q_len, self.num_heads, self.head_dim + ).transpose(1, 2) + key_states = key_states.view( + bsz, q_len, self.num_query_groups, self.head_dim + ).transpose(1, 2) + value_states = value_states.view( + bsz, q_len, self.num_query_groups, self.head_dim + ).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if layer_past is not None: + kv_seq_len += layer_past[0].shape[-2] + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin, position_ids + ) + + if layer_past is not None: + # reuse k, v, self_attention + key_states = torch.cat([layer_past[0], key_states], dim=2) + value_states = torch.cat([layer_past[1], value_states], dim=2) + + layer_past = (key_states, value_states) if use_cache else None + using_pytorch = False + using_ixformer = True + + if using_pytorch: + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_heads // self.kv_heads) + value_states = repeat_kv(value_states, self.num_heads // self.kv_heads) + attn_weights = torch.matmul( + query_states, key_states.transpose(2, 3) + ) / math.sqrt(self.head_dim) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + mask_value = self._get_mask_value( + attn_weights.device, attn_weights.dtype + ) + # The fused kernel is very slow when the key length is not a multiple of 8, so we skip fusion. + attn_weights = torch.where(attention_mask, attn_weights, mask_value) + + # upcast attention to fp32 + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(query_states.dtype) + attn_weights = self.attn_dropout(attn_weights) + attn_output = torch.matmul(attn_weights, value_states) + if using_ixformer: + # print(f"query_states.dtype {query_states.dtype, key_states.dtype,value_states.dtype}") + attn_output = mha(query_states, key_states, value_states, attention_mask) + # print(attn_output.shape,attn_output_ixformer.shape) + # diff = attn_output-attn_output_ixformer + # print(f"attention_mask {attention_mask.dtype,attention_mask.shape}") + # print(f"diff {diff.max()} attention_mask is None {attention_mask is None}") + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.embed_dim) + + attn_output = self.c_proj(attn_output) + attn_output = self.resid_dropout(attn_output) + + outputs = (attn_output, layer_past) + if output_attentions: + outputs += (attn_weights,) + + return outputs # a, present, (attentions) + + +class CodeShellMLP(nn.Module): + def __init__(self, intermediate_size, config): + super().__init__() + embed_dim = config.hidden_size + self.c_fc = nn.Linear(embed_dim, intermediate_size) + self.c_proj = nn.Linear(intermediate_size, embed_dim) + self.act = ACT2FN[config.activation_function] + self.dropout = nn.Dropout(config.resid_pdrop) + + # Copied from transformers.models.gpt2.modeling_gpt2.GPT2MLP.forward + def forward(self, hidden_states: Optional[Tuple[torch.Tensor]]) -> torch.Tensor: + using_pytorch = False + using_ixformer = True + # hidden_states_ixf = hidden_states.clone() + if using_pytorch: + hidden_states = self.c_fc(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.c_proj(hidden_states) + hidden_states = self.dropout(hidden_states) + if using_ixformer: + hidden_states = mlp_forward(self, hidden_states) + # diff = (hidden_states-hidden_states_ixf) + # print(f"mlp diff {diff.max()}") + return hidden_states + + +class CodeShellBlock(nn.Module): + def __init__(self, config, layer_idx=None): + super().__init__() + hidden_size = config.hidden_size + self.inner_dim = ( + config.n_inner if config.n_inner is not None else 4 * hidden_size + ) + + self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + self.attn = CodeShellAttention(config, layer_idx=layer_idx) + self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + + self.mlp = CodeShellMLP(self.inner_dim, config) + + def forward( + self, + hidden_states: Optional[Tuple[torch.Tensor]], + layer_past: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + ) -> Union[ + Tuple[torch.Tensor], + Tuple[torch.Tensor, torch.Tensor], + Tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ]: + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attn_outputs = self.attn( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask, + use_cache=use_cache, + output_attentions=output_attentions, + ) + attn_output = attn_outputs[0] # output_attn: a, present, (attentions) + + outputs = attn_outputs[1:] + # residual connection + hidden_states = attn_output + residual + + residual = hidden_states + hidden_states = self.ln_2(hidden_states) + feed_forward_hidden_states = self.mlp(hidden_states) + # residual connection + hidden_states = residual + feed_forward_hidden_states + + if use_cache: + outputs = (hidden_states,) + outputs + else: + outputs = (hidden_states,) + outputs[1:] + + return outputs # hidden_states, present, (attentions, cross_attentions) + + +class CodeShellPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = CodeShellConfig + base_model_prefix = "transformer" + supports_gradient_checkpointing = True + _no_split_modules = ["ShellBlock"] + _skip_keys_device_placement = "past_key_values" + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights(self, module): + """Initialize the weights.""" + if isinstance(module, (CodeShellMLP, CodeShellAttention)): + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + module.c_proj.weight.data.normal_( + mean=0.0, + std=( + self.config.initializer_range / math.sqrt(2 * self.config.n_layer) + ), + ) + module.c_proj._is_hf_initialized = True + elif isinstance(module, nn.Linear): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + # Copied from transformers.models.gpt2.modeling_gpt2.GPT2PreTrainedModel._set_gradient_checkpointing with GPT2->Shell + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, CodeShellModel): + module.gradient_checkpointing = value + + +GPT_BIGCODE_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + Parameters: + config ([`CodeShellConfig`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +GPT_BIGCODE_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`): + `input_ids_length` = `sequence_length` if `past_key_values` is `None` else + `past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input + sequence tokens in the vocabulary. + If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as + `input_ids`. + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + [What are input IDs?](../glossary#input-ids) + past_key_values (`Tuple[torch.Tensor]` of length `config.n_layers`): + Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see + `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have + their past given to this model should not be passed as `input_ids` as they have already been computed. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + If `past_key_values` is used, `attention_mask` needs to contain the masking strategy that was used for + `past_key_values`. In other words, the `attention_mask` always has to have the length: + `len(past_key_values) + len(input_ids)` + [What are attention masks?](../glossary#attention-mask) + token_type_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`, *optional*): + Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, + 1]`: + - 0 corresponds to a *sentence A* token, + - 1 corresponds to a *sentence B* token. + [What are token type IDs?](../glossary#token-type-ids) + position_ids (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + [What are position IDs?](../glossary#position-ids) + head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see + `past_key_values`). + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The bare GPT_BIGCODE Model transformer outputting raw hidden-states without any specific head on top.", + GPT_BIGCODE_START_DOCSTRING, +) +class CodeShellModel(CodeShellPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.group_query_attention = config.group_query_attention + self.num_query_groups = config.num_query_groups + self.position_embedding_type = config.position_embedding_type + self.embed_dim = config.hidden_size + + self.wte = nn.Embedding(config.vocab_size, self.embed_dim) + if self.position_embedding_type == "learned_absolute": + self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim) + else: + pass + + self.drop = nn.Dropout(config.embd_pdrop) + self.h = nn.ModuleList( + [ + CodeShellBlock(config, layer_idx=i) + for i in range(config.num_hidden_layers) + ] + ) + self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) + + max_positions = config.max_position_embeddings + self.register_buffer( + "bias", + torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)), + persistent=False, + ) + + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.wte + + def set_input_embeddings(self, new_embeddings): + self.wte = new_embeddings + + @add_start_docstrings_to_model_forward(GPT_BIGCODE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.Tensor]] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time" + ) + elif input_ids is not None: + input_shape = input_ids.size() + input_ids = input_ids.reshape(-1, input_shape[-1]) + batch_size = input_ids.shape[0] + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + batch_size = inputs_embeds.shape[0] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if batch_size <= 0: + raise ValueError("batch_size has to be defined and > 0") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if token_type_ids is not None: + token_type_ids = token_type_ids.reshape(-1, input_shape[-1]) + if position_ids is not None: + position_ids = position_ids.reshape(-1, input_shape[-1]) + + if past_key_values is None: + past_length = 0 + past_key_values = tuple([None] * len(self.h)) + else: + past_length = past_key_values[0][0].size(-2) + + if ( + attention_mask is not None + and len(attention_mask.shape) == 2 + and position_ids is None + ): + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_length > 0: + position_ids = position_ids[ + :, past_length : input_shape[-1] + past_length : + ] + elif position_ids is None: + position_ids = torch.arange( + past_length, + input_shape[-1] + past_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0).reshape(-1, input_shape[-1]) + + # Self-attention mask. + query_length = input_shape[-1] + key_length = past_length + query_length + self_attention_mask = self.bias[ + None, key_length - query_length : key_length, :key_length + ] + + if attention_mask is not None: + self_attention_mask = self_attention_mask * attention_mask.reshape( + batch_size, 1, -1 + ).to(dtype=torch.bool, device=self_attention_mask.device) + + # MQA models: (batch_size, query_length, n_heads, key_length) + # MHA models: (batch_size, n_heads, query_length, key_length) + attention_mask = self_attention_mask.unsqueeze(1) + + encoder_attention_mask = None + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # head_mask has shape n_layer x batch x n_heads x N x N + head_mask = self.get_head_mask(head_mask, self.config.n_layer) + + if inputs_embeds is None: + inputs_embeds = self.wte(input_ids) + + hidden_states = inputs_embeds + if self.position_embedding_type == "learned_absolute": + position_embeds = self.wpe(position_ids) + hidden_states = hidden_states + position_embeds + + if token_type_ids is not None: + token_type_embeds = self.wte(token_type_ids) + hidden_states = hidden_states + token_type_embeds + + hidden_states = self.drop(hidden_states) + + output_shape = input_shape + (hidden_states.size(-1),) + + presents = [] if use_cache else None + all_self_attentions = () if output_attentions else None + all_hidden_states = () if output_hidden_states else None + for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + # None for past_key_value + return module(*inputs, use_cache, output_attentions) + + return custom_forward + + outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + None, + attention_mask, + position_ids, + head_mask[i], + encoder_hidden_states, + encoder_attention_mask, + ) + else: + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask[i], + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + ) + + hidden_states = outputs[0] + if use_cache: + presents.append(outputs[1]) + + if output_attentions: + all_self_attentions = all_self_attentions + ( + outputs[2 if use_cache else 1], + ) + + hidden_states = self.ln_f(hidden_states) + hidden_states = hidden_states.reshape(output_shape) + # Add last hidden state + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + presents, + all_hidden_states, + all_self_attentions, + ] + if v is not None + ) + + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=presents, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +class EndOfFunctionCriteria(StoppingCriteria): + """Custom `StoppingCriteria` which checks if all generated functions in the batch are completed.""" + + def __init__(self, input_lengths, eof_strings, tokenizer): + self.input_lengths = input_lengths + self.eof_strings = eof_strings + self.tokenizer = tokenizer + + def __call__(self, input_ids, scores, **kwargs): + """Returns true if all generated sequences contain any of the end-of-function strings.""" + decoded_generations = [] + for _input_ids, input_length in zip(input_ids, self.input_lengths): + decoded_generations.append(self.tokenizer.decode(_input_ids[input_length:])) + done = [] + for decoded_generation in decoded_generations: + done.append( + any( + [ + stop_string in decoded_generation + for stop_string in self.eof_strings + ] + ) + ) + return all(done) + + +class TextIterStreamer: + def __init__(self, tokenizer, skip_prompt=False, skip_special_tokens=False): + self.tokenizer = tokenizer + self.skip_prompt = skip_prompt + self.skip_special_tokens = skip_special_tokens + self.tokens = [] + self.text_queue = Queue() + self.next_tokens_are_prompt = True + + def put(self, value): + if self.skip_prompt and self.next_tokens_are_prompt: + self.next_tokens_are_prompt = False + else: + if len(value.shape) > 1: + value = value[0] + self.tokens.extend(value.tolist()) + self.text_queue.put( + self.tokenizer.decode( + self.tokens, skip_special_tokens=self.skip_special_tokens + ) + ) + + def end(self): + self.text_queue.put(None) + + def __iter__(self): + return self + + def __next__(self): + value = self.text_queue.get() + if value is None: + raise StopIteration() + else: + return value + + +@add_start_docstrings( + """ + The GPT_BIGCODE Model transformer with a language modeling head on top (linear layer with weights tied to the input + embeddings). + """, + GPT_BIGCODE_START_DOCSTRING, +) +class CodeShellForCausalLM(CodeShellPreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.transformer = CodeShellModel(config) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def quantize(self, bits: int): + try: + import bitsandbytes + + from .quantizer import quantize + except ImportError: + raise ImportError(f"Needs bitsandbytes to run quantize.") + return quantize(self, bits) + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def prepare_inputs_for_generation( + self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs + ): + token_type_ids = kwargs.get("token_type_ids", None) + # only last token for inputs_ids if past is defined in kwargs + if past_key_values: + input_ids = input_ids[:, -1].unsqueeze(-1) + if token_type_ids is not None: + token_type_ids = token_type_ids[:, -1].unsqueeze(-1) + + attention_mask = kwargs.get("attention_mask", None) + position_ids = kwargs.get("position_ids", None) + + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -1].unsqueeze(-1) + else: + position_ids = None + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "position_ids": position_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + } + ) + return model_inputs + + @add_start_docstrings_to_model_forward(GPT_BIGCODE_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]: + r""" + labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set + `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` + are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + lm_logits = self.lm_head(hidden_states) + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = lm_logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous().to(shift_logits.device) + # Flatten the tokens + loss_fct = CrossEntropyLoss() + loss = loss_fct( + shift_logits.reshape(-1, shift_logits.size(-1)), + shift_labels.reshape(-1), + ) + + if not return_dict: + output = (lm_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutputWithCrossAttentions( + loss=loss, + logits=lm_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += ( + tuple( + past_state.index_select(0, beam_idx.to(past_state.device)) + for past_state in layer_past + ), + ) + return reordered_past + + def build_chat_input(self, query, history, tokenizer, max_new_tokens=None): + user_name = "## human:" + ai_name = "## assistant: " + stop = "||" + + prompt = "" + for q, r in history: + prompt += f"{user_name}{q}{stop}" + prompt += f"{ai_name}{r}{stop}" + prompt += f"{user_name}{query}{stop}" + prompt += ai_name.rstrip() + + max_new_tokens = max_new_tokens or self.generation_config.max_new_tokens + max_new_tokens = max_new_tokens or 128 + max_input_tokens = self.config.n_positions - max_new_tokens + + input_tokens = tokenizer.encode(prompt) + input_tokens = input_tokens[-max_input_tokens:] # truncate left + return torch.LongTensor([input_tokens]).to(self.device) + + def chat( + self, + query, + history, + tokenizer, + stream=False, + generation_config: Optional[GenerationConfig] = None, + ): + generation_config = generation_config or self.generation_config + input_ids = self.build_chat_input( + query, history, tokenizer, generation_config.max_new_tokens + ) + stopping_criteria = StoppingCriteriaList( + [ + EndOfFunctionCriteria( + [len(input_ids[0])], + ["||", "|end|", "<|endoftext|>", "## human"], + tokenizer, + ) + ] + ) + + if stream: + streamer = TextIterStreamer( + tokenizer, skip_prompt=True, skip_special_tokens=True + ) + Thread( + target=self.generate, + kwargs=dict( + inputs=input_ids, + streamer=streamer, + stopping_criteria=stopping_criteria, + generation_config=generation_config, + ), + ).start() + return streamer + else: + outputs = self.generate( + input_ids, + generation_config=generation_config, + stopping_criteria=stopping_criteria, + ) + response = tokenizer.decode( + outputs[0][len(input_ids[0]) :], skip_special_tokens=True + ) + return response + + def generate_stream(self, prompt, tokenizer, generation_config=None, **kwargs): + generation_config = generation_config or self.generation_config + max_input_tokens = ( + self.config.n_positions - self.generation_config.max_new_tokens + ) + + input_ids = tokenizer.encode(prompt) + input_ids = input_ids[-max_input_tokens:] # truncate left + + stopping_criteria = StoppingCriteriaList( + [ + EndOfFunctionCriteria( + [len(input_ids[0])], + ["||", "|end|", "<|endoftext|>", "## human"], + tokenizer, + ) + ] + ) + + streamer = TextIterStreamer( + tokenizer, skip_prompt=True, skip_special_tokens=True + ) + Thread( + target=self.generate, + kwargs=dict( + inputs=input_ids, stopping_criteria=stopping_criteria, **kwargs + ), + ).start() + return streamer + + +class CodeShell4bitForCausalLM(CodeShellForCausalLM): + def __init__(self, config): + CodeShellPreTrainedModel.__init__(self, config) + self.transformer = CodeShellModel(config) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + + try: + import bitsandbytes + + from .quantizer import quantize_offline + + quantize_offline(self) + except ImportError: + raise ImportError(f"Needs bitsandbytes to run quantize.") + + self.post_init() + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + *model_args, + config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None, + cache_dir: Optional[Union[str, os.PathLike]] = None, + ignore_mismatched_sizes: bool = False, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + use_safetensors: bool = None, + **kwargs, + ): + if not isinstance(config, PretrainedConfig): + config_path = ( + config if config is not None else pretrained_model_name_or_path + ) + config, _ = cls.config_class.from_pretrained( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + force_download=force_download, + resume_download=False, + proxies=None, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder="", + _from_auto=False, + _from_pipeline=None, + **kwargs, + ) + + # Load config if we don't provide a configuration + from .quantizer import load_state_dict_for_qunantied_model + + model = cls(config) + state_dict = torch.load( + os.path.join(pretrained_model_name_or_path, "pytorch_model.bin"), + map_location="cpu", + ) + model = load_state_dict_for_qunantied_model(model, state_dict) + model.eval() + + # If it is a model with generation capabilities, attempt to load the generation config + if model.can_generate(): + try: + model.generation_config = GenerationConfig.from_pretrained( + pretrained_model_name_or_path, + cache_dir=cache_dir, + force_download=force_download, + resume_download=False, + proxies=None, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder="", + _from_auto=False, + _from_pipeline=None, + **kwargs, + ) + except (OSError, TypeError): + pass + + device_map = kwargs.pop("device_map", None) + if device_map is not None: + model = model.to(torch.device(device_map)) + + return model diff --git a/ixformer_sdk/inference/models/codeshell_7b_chat/modeling_codeshell_ixformer.py b/ixformer_sdk/inference/models/codeshell_7b_chat/modeling_codeshell_ixformer.py new file mode 100644 index 00000000..cb3538ee --- /dev/null +++ b/ixformer_sdk/inference/models/codeshell_7b_chat/modeling_codeshell_ixformer.py @@ -0,0 +1,99 @@ +import math + +import ixformer.functions as ixf_F +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +def mha(query, key, value, attention_mask): + if attention_mask is None and query.shape[2] == key.shape[2]: + context_layer = ixf_F.scaled_dot_product_attention( + query.contiguous(), key.contiguous(), value.contiguous(), is_causal=True + ) + else: + # if attention_mask is not None: + # # attention_mask = attention_mask + # attention_mask = (~attention_mask).cuda().float()*(-10000) + context_layer = ixf_F.scaled_dot_product_attention( + query.contiguous(), key.contiguous(), value.contiguous(), attention_mask + ) + + # context_layer = context_layer.transpose(1, 2).contiguous() + # res_shape = list(context_layer.shape) + # res_shape = res_shape[:2] + [-1] + # context_layer = context_layer.view(*res_shape) + return context_layer + # batch_size, head_num, seq_len, head_dim = query.shape + # src_len = query.shape[-2] + # tgt_len = key.shape[-2] + + # if attention_mask is None and src_len == tgt_len: + # attention_mask = ~torch.tril(torch.ones([src_len, tgt_len])).bool() + # elif attention_mask is None: + # attention_mask = torch.zeros([src_len, tgt_len]) + # attention_mask = attention_mask.cuda().int() + + # attention_scores = ixf_F.act_bias_mm( + # query, key, scale=1 / math.sqrt(head_dim), trans_format="TN" + # ) + # # softmax + # # if tgt_len > 2048: + # # if not (attention_mask == 0).all(): + # # attention_scores.masked_fill_(attention_mask.bool(), -10000.0) + # # dtype = attention_scores.dtype + # # attention_probs = F.softmax(attention_scores.float(), dim=-1) + # # attention_probs = attention_probs.type(dtype) + # # else: + # # raise NotImplementedError() + # attention_probs = ixf_F.attention_masked_softmax( + # attention_scores, attention_mask.int() + # ) + # # s * v + # # batch_size,head_num,seq_len,head_dim + # context_layer = ixf_F.act_bias_mm( + # attention_probs, value, trans_format="NN") + # context_layer = context_layer.transpose(1, 2).contiguous() + # context_layer = context_layer.view( + # batch_size, seq_len, head_num * head_dim) + + +def mlp(mlp_input, ff1_weight, ff1_bias, ff2_weight): + input_shape = list(mlp_input.shape) + mlp_input = mlp_input.view(-1, input_shape[-1]) + mlp_output = ixf_F.act_bias_mm( + mlp_input, ff1_weight, ff1_bias, scale=1, act_type="gelu", trans_format="TN" + ) + mlp_output = ixf_F.linear(mlp_output, ff2_weight, None) + input_shape[-1] = -1 + mlp_output = mlp_output.view(*input_shape) + return mlp_output + + +def mlp_forward(self, hidden_states): + # [s, b, 4hp] + # intermediate_parallel = self.dense_h_to_4h(hidden_states) + # intermediate_parallel = self.activation_func(intermediate_parallel) + input_shape = list(hidden_states.shape) + hidden_states = hidden_states.view(-1, input_shape[-1]) + mlp_output = ixf_F.act_bias_mm( + hidden_states, + self.c_fc.weight, + self.c_fc.bias, + scale=1, + act_type="gelu", + trans_format="TN", + ) + if isinstance(self.c_proj, nn.Linear): + output = ixf_F.linear( + mlp_output, + self.c_proj.weight, + self.c_proj.bias, + ) + else: + output = self.c_proj(mlp_output) + output = output.view(*input_shape) + return output diff --git a/ixformer_sdk/inference/overlap/__init__.py b/ixformer_sdk/inference/overlap/__init__.py new file mode 100644 index 00000000..8b6933bd --- /dev/null +++ b/ixformer_sdk/inference/overlap/__init__.py @@ -0,0 +1 @@ +from .llama_decoder_layer_overlap import LlamaDecoderLayerOverlapProtocol, LlamaDecoderLayerOverlapDefault, create_vllm_llama_decoder_layer \ No newline at end of file diff --git a/ixformer_sdk/inference/overlap/fmha_oproj_allreduce_ln_gating_overlap.py b/ixformer_sdk/inference/overlap/fmha_oproj_allreduce_ln_gating_overlap.py new file mode 100644 index 00000000..5524cd64 --- /dev/null +++ b/ixformer_sdk/inference/overlap/fmha_oproj_allreduce_ln_gating_overlap.py @@ -0,0 +1,333 @@ +import dataclasses +from typing import Optional, Tuple + +import ixformer.distributed as ixfd +import ixformer.functions as F +import torch +import torch.distributed as dist +from ixformer.contrib.vllm_flash_attn import flash_attn_varlen_func +from ixformer.distributed.overlap_comm import SplitOverlapComm + +from ixformer.core import config as ixff_config + + +@dataclasses.dataclass +class FmhaOProjAllReduceLnGatingParams: + # ============================== + # attention + # ============================== + + # shape: [Batch * SeqLen, NumHeads / TP, HeadDim] + q: torch.Tensor + + # shape: [Batch * SeqLen, NumHeads / TP, HeadDim] + k: torch.Tensor + + # shape: [Batch * SeqLen, NumHeads / TP, HeadDim] + v: torch.Tensor + + # shape [Batch + 1], dtype torch.int32. The cumulative sequence lengths + # of the sequences in the batch, used to index into q. + cu_seqlens_q: torch.Tensor + + # shape: [Batch + 1], dtype torch.int32. The cumulative sequence lengths + # of the sequences in the batch, used to index into kv. + cu_seqlens_k: torch.Tensor + + # Maximum query sequence length in the batch. + max_seqlen_q: int + + # Maximum key sequence length in the batch. + max_seqlen_k: int + + # ============================== + # o_proj + # ============================== + + # shape: [HiddenSize, NumHeads * HeadDim / TP], dtype: int8 + o_proj_weight: torch.Tensor + + # shape: [HiddenSize], dtype: float32 + o_proj_weight_scale: torch.Tensor + + # shape: [HiddenSize] + o_proj_bias: torch.Tensor + + # shape: [NumHeads * HeadDim / TP], dtype: float16 or bfloat16 + o_proj_smooth_scale: torch.Tensor + + # ============================== + # ln + # ============================== + + # shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16 + residual: torch.Tensor + + # shape: [HiddenSize], dtype: float16 or bfloat16 + ln_weight: torch.Tensor + + # shape: [HiddenSize], dtype: float16 or bfloat16 + ln_bias: torch.Tensor + + # ============================== + # gating linear + # ============================== + + # shape: [TopK, HiddenSize], dtype: float16 or bfloat16 + gating_weight: torch.Tensor + + # shape: [SeqLen, TopK], dtype: float16 or bfloat16 + out: Optional[torch.Tensor] = None + + # ============================== + # default parameters + # ============================== + + # the seqlens of q for per chunk when using overlap, + # the parameter can be initiated by params.prepare_overlap_params(), + # and only need to initialize once during the model's forward. + cu_seqlens_q_chunks = None + cu_seqlens_k_chunks = None + + softmax_scale: Optional[float] = None + ln_eps: float = 1e-5 + + @property + def batch(self): + return len(self.cu_seqlens_q) - 1 + + @property + def seqlen(self): + return self.q.shape[0] + + @property + def topk(self): + return self.gating_weight.shape[0] + + def prepare_overlap_params(self): + """compute the cu_seqlens qk of chunk when using overlap""" + first_chunk_size = int(self.q.shape[0] // 2) + + if not hasattr(self.cu_seqlens_q, "q_chunks"): + first_q_chunks_cu_seqlens = self.cu_seqlens_q.clone() + first_q_chunks_cu_seqlens[-1] = first_chunk_size + + self.cu_seqlens_q.q_chunks = [ + first_q_chunks_cu_seqlens, + first_q_chunks_cu_seqlens, + ] + + self.cu_seqlens_q_chunks = self.cu_seqlens_q.q_chunks + + if not hasattr(self.cu_seqlens_k, "k_chunks"): + first_chunk = self.cu_seqlens_k.clone() + last_chunk = self.cu_seqlens_k + first_chunk[-1] = first_chunk_size + self.cu_seqlens_k.k_chunks = [first_chunk, last_chunk] + + self.cu_seqlens_k_chunks = self.cu_seqlens_k.k_chunks + + return self + + +class FmhaOProjAllreduceLnGatingOverlap(SplitOverlapComm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + if self.num_chunks != 2: + raise RuntimeError( + f"Overlap only support num_chunks == 2, but got {self.num_chunks}." + ) + + self.allreduce_end_events = [torch.cuda.Event() for _ in range(self.num_chunks)] + + def start_ln_gating(self, chunk_idx): + compute_stream = self._compute_streams[chunk_idx % self.num_compute_streams] + compute_stream.wait_event(self.allreduce_end_events[chunk_idx]) + + def compute(self, params: FmhaOProjAllReduceLnGatingParams): + if params.out is None: + params.out = torch.empty( + [params.seqlen, params.topk], device="cuda", dtype=torch.float + ) + + seqlen_chunks = [int(params.q.shape[0] // 2)] + seqlen_chunks.append(params.q.shape[0] - seqlen_chunks[0]) + + q_chunks = torch.split_with_sizes(params.q, seqlen_chunks, dim=0) + + ar_out_chunks = [] + for chunk_idx in range(len(seqlen_chunks)): + with self.compute_stream_context(chunk_idx): + hidden_states = flash_attn_varlen_func( + q=q_chunks[chunk_idx], + k=params.k, + v=params.v, + cu_seqlens_q=params.cu_seqlens_q_chunks[chunk_idx], + cu_seqlens_k=params.cu_seqlens_k_chunks[chunk_idx], + max_seqlen_q=seqlen_chunks[chunk_idx], + max_seqlen_k=seqlen_chunks[0] + if chunk_idx == 0 + else params.max_seqlen_k, + softmax_scale=params.softmax_scale, + causal=True, + window_size=(-1, -1), + alibi_slopes=None, + softcap=0, + ) + + hidden_states = hidden_states.view(hidden_states.shape[0], -1) + hidden_states, i_scales = F.dynamic_scaled_quant_dynamic_int8( + hidden_states, params.o_proj_smooth_scale + ) + + out_chunk = F.w8a8( + hidden_states, + params.o_proj_weight, + i_scales, + params.o_proj_weight_scale, + bias=params.o_proj_bias, + out_dtype=params.residual.dtype, + output=None, + persistent=True, + ) + + self.start_comm(chunk_idx) + + ixfd.all_reduce( + out_chunk, async_op=True, group=self.comm_group, use_comm_stream=True + ) + ar_out_chunks.append(out_chunk) + + self.allreduce_end_events[chunk_idx].record(self._comm_stream) + + ln_out = torch.empty_like(params.residual) + ln_out_chunks = ln_out.chunk(2, dim=0) + residual_chunks = torch.split_with_sizes(params.residual, seqlen_chunks, dim=0) + + if params.out is None: + params.out = torch.empty( + [params.seqlen, params.topk], dtype=params.residual.dtype, device="cuda" + ) + + out_chunks = list(torch.split_with_sizes(params.out, seqlen_chunks, dim=0)) + + for chunk_idx in range(len(seqlen_chunks)): + self.start_ln_gating(chunk_idx) + with self.compute_stream_context(chunk_idx): + ln_out_chunk, residual_chunk = F.residual_layer_norm( + input=ar_out_chunks[chunk_idx], + weight=params.ln_weight, + bias=params.ln_bias, + residual=residual_chunks[chunk_idx].reshape( + ar_out_chunks[chunk_idx].shape + ), + eps=params.ln_eps, + output=ln_out_chunks[chunk_idx], + ) + if ln_out_chunk.dtype == params.gating_weight.dtype: + F.linear( + ln_out_chunk, params.gating_weight, output=out_chunks[chunk_idx] + ) + else: + F.mixed_type_linear( + ln_out_chunk, params.gating_weight, output=out_chunks[chunk_idx] + ) + + return ( + params.residual.reshape(params.batch, params.seqlen, -1), + ln_out.reshape(params.batch, params.seqlen, -1), + params.out, + ) + + +_fa_o_proj_allreduce_ln_gating_overlap = None + + +def fmha_oproj_allreduce_ln_gating( + params: FmhaOProjAllReduceLnGatingParams, + enable_overlap: bool = False, + comm_group=None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + FMHA + OProjLinear + AllReduce + LayerNorm + GatingLinear + + Args: + params: fused operator params + enable_overlap: whether enable overlap + comm_group: communication group + Returns: + Residual: shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16 + HiddenStates: shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16 + GatingLinearOutput: shape: [SeqLen, TopK], dtype: float16 or bfloat16 + """ + + global _fa_o_proj_allreduce_ln_gating_overlap + if _fa_o_proj_allreduce_ln_gating_overlap is None: + _fa_o_proj_allreduce_ln_gating_overlap = ( + FmhaOProjAllreduceLnGatingOverlap.dispatcher( + num_chunks=2, comm_group=comm_group + ).forward + ) + + if ( + enable_overlap + and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM + and dist.is_initialized() + and dist.get_world_size(comm_group) > 1 + and params.batch == 1 + and params.seqlen > 1 + ): + if params.cu_seqlens_q_chunks is None: + params = params.prepare_overlap_params() + return _fa_o_proj_allreduce_ln_gating_overlap(params) + + hidden_states = flash_attn_varlen_func( + q=params.q, + k=params.k, + v=params.v, + cu_seqlens_q=params.cu_seqlens_q, + cu_seqlens_k=params.cu_seqlens_k, + max_seqlen_q=params.max_seqlen_q, + max_seqlen_k=params.max_seqlen_k, + softmax_scale=params.softmax_scale, + causal=True, + window_size=(-1, -1), + alibi_slopes=None, + softcap=0, + ) + + input = hidden_states.view(hidden_states.shape[0], -1) + input, i_scales = F.dynamic_scaled_quant_smoothquant( + input, params.o_proj_smooth_scale + ) + + hidden_states = F.w8a8( + input, + params.o_proj_weight, + i_scales, + params.o_proj_weight_scale, + bias=params.o_proj_bias, + out_dtype=params.residual.dtype, + output=None, + ) + ixfd.all_reduce(hidden_states, async_op=True, group=comm_group) + + hidden_states, residual = F.residual_layer_norm( + input=hidden_states, + weight=params.ln_weight, + bias=params.ln_bias, + residual=params.residual.reshape(hidden_states.shape), + eps=params.ln_eps, + ) + if hidden_states.dtype == params.gating_weight.dtype: + out = F.linear(hidden_states, params.gating_weight, output=params.out) + else: + out = F.mixed_type_linear( + hidden_states, params.gating_weight, output=params.out + ) + return ( + residual.reshape(params.batch, params.seqlen, -1), + hidden_states.reshape(params.batch, params.seqlen, -1), + out, + ) diff --git a/ixformer_sdk/inference/overlap/group_gemm_moe_reduce_sum_allreduce_overlap.py b/ixformer_sdk/inference/overlap/group_gemm_moe_reduce_sum_allreduce_overlap.py new file mode 100644 index 00000000..ceaa20cb --- /dev/null +++ b/ixformer_sdk/inference/overlap/group_gemm_moe_reduce_sum_allreduce_overlap.py @@ -0,0 +1,219 @@ +import dataclasses +import math +from contextlib import contextmanager +from typing import List, Optional + +import ixformer.distributed as ixfd +import ixformer.functions as F +import torch +import torch.distributed as dist +from ixformer.distributed.overlap_comm import SplitOverlapComm + +from ixformer.core import config as ixff_config + + +@dataclasses.dataclass +class GroupGemmMoeReduceSumAllReduceParams: + # M: NumTokens * TopK + # K: InnerSize // TP + # N: HiddenSize + # NumTokens: M // TopK + + # ================================= + # group gemm + # ================================= + + # the top k of experts + topk: int + + # shape: [M, K] if format[1]=="N" else [K, M], dtype: int8 + input: torch.Tensor + + # shape: [NumExperts, N, K] if format[0]=="T" else [NumExperts, K, N], dtype: int8 + weight: torch.Tensor + + # shape: [M], dtype: float32 + i_scales: torch.Tensor + + # shape: [NumExperts, N], dtype: float32 + w_scales: torch.Tensor + + # shape: [NumExperts], dtype: int32 + tokens_per_experts: torch.Tensor + + # the dtype of output, support float16 and bfloat16 + out_dtype: torch.dtype = None + + # index of dst to src, shape: [M], dtype: int32 + dst_to_src: torch.Tensor = None + + # only support TN now + format: str = "TN" + + # ================================= + # moe reduce sum + # ================================= + + # shape: [M // TopK, TopK], dtype: torch.float16 or torch.bfloat16 + topk_weight: torch.Tensor = None + + # shape: [M // TopK, N], dtype: torch.float16 or torch.bfloat16 + output: torch.Tensor = None + + # overlap + output_chunks: Optional[List[torch.Tensor]] = None + + @property + def M(self): + return self.input.shape[0] + + @property + def N(self): + if torch.is_tensor(self.weight): + return self.weight.shape[1] + return sum(t.shape[1] for t in self.weight) + + @property + def K(self): + return self.input.shape[-1] + + def prepare_overlap_params( + self, num_chunks: int, split_ratio: Optional[float] = None + ): + if num_chunks == 2 and split_ratio not in [0, None]: + return self.prepare_overla_params_with_ratio(split_ratio) + return self.prepare_overlap_params_with_chunks(num_chunks) + + def prepare_overlap_params_with_chunks(self, num_chunks: int): + if torch.is_tensor(self.weight): + weight_chunks = torch.chunk(self.weight, num_chunks, dim=1) + self.weight = list(weight_chunks) + + if torch.is_tensor(self.w_scales): + weight_scale_chunks = torch.chunk(self.w_scales, num_chunks, dim=1) + self.w_scales = list(weight_scale_chunks) + + if self.output is None: + self.output = torch.empty( + self.M // self.topk, self.N, dtype=self.out_dtype, device="cuda" + ) + + if torch.is_tensor(self.output): + output_chunks = torch.chunk(self.output, num_chunks, dim=1) + self.output_chunks = list(output_chunks) + + def prepare_overla_params_with_ratio(self, split_ratio: float): + N = self.N + n_chunks = [int(math.ceil(N * split_ratio))] + n_chunks.append(N - n_chunks[0]) + + if torch.is_tensor(self.weight): + weight_chunks = torch.split(self.weight, n_chunks, dim=1) + self.weight = list(weight_chunks) + + if torch.is_tensor(self.w_scales): + weight_scale_chunks = torch.split(self.w_scales, n_chunks, dim=1) + self.w_scales = list(weight_scale_chunks) + + if self.output is None: + self.output = torch.empty( + self.M // self.topk, self.N, dtype=self.out_dtype, device="cuda" + ) + + if torch.is_tensor(self.output): + output_chunks = torch.split(self.output, n_chunks, dim=1) + self.output_chunks = list(output_chunks) + + +class GroupGemmMoeReduceSumAllReduceSplitNOverlap(SplitOverlapComm): + def compute(self, params: GroupGemmMoeReduceSumAllReduceParams): + for chunk_idx, (weight, weight_scale) in enumerate( + zip(params.weight, params.w_scales) + ): + with self.compute_stream_context(chunk_idx): + out = F.moe_w8a8_group_gemm( + input=params.input, + weight=weight, + i_scales=params.i_scales, + w_scales=weight_scale, + output_dtype=params.out_dtype, + tokens_per_experts=params.tokens_per_experts, + dst_to_src=params.dst_to_src, + format=params.format, + ) + out = out.reshape(-1, params.topk, out.shape[-1]) + out = F.moe_output_reduce_sum( + input=out, + topk_weight=params.topk_weight, + output=params.output_chunks[chunk_idx], + ) + + self.start_comm(chunk_idx) + ixfd.all_reduce( + out, + async_op=True, + group=self.comm_group, + use_comm_stream=True, + algo=ixfd.AllReduceAlgo.Stride, + ) + + # if chunk_idx == 0: torch.cuda.synchronize() + + return params.output + + +_group_gemm_moe_reduce_sum_all_reduce_overlap = None + + +def group_gemm_moe_reduce_sum_allreduce( + params: GroupGemmMoeReduceSumAllReduceParams, + enable_overlap: bool = False, + comm_group=None, + num_chunks=2, + split_ratio: Optional[float] = None, +): + if params.output is None and params.out_dtype is None: + raise RuntimeError( + "group_gemm_moe_reduce_sum_all_reduce need out_dtype argument when output is none." + ) + + if params.out_dtype is None: + params.out_dtype = params.output.dtype + + if ( + enable_overlap + and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM + and dist.is_initialized() + and dist.get_world_size(comm_group) > 1 + ): + global _group_gemm_moe_reduce_sum_all_reduce_overlap + if _group_gemm_moe_reduce_sum_all_reduce_overlap is None: + _group_gemm_moe_reduce_sum_all_reduce_overlap = ( + GroupGemmMoeReduceSumAllReduceSplitNOverlap.dispatcher( + num_chunks=num_chunks, comm_group=comm_group + ).forward + ) + + params.prepare_overlap_params(num_chunks=num_chunks, split_ratio=split_ratio) + return _group_gemm_moe_reduce_sum_all_reduce_overlap(params) + + out = F.moe_w8a8_group_gemm( + input=params.input, + weight=params.weight, + i_scales=params.i_scales, + w_scales=params.w_scales, + output_dtype=params.out_dtype, + tokens_per_experts=params.tokens_per_experts, + dst_to_src=params.dst_to_src, + format=params.format, + ) + + out = out.reshape(-1, params.topk, out.shape[-1]) + out = F.moe_output_reduce_sum( + input=out, topk_weight=params.topk_weight, output=params.output + ) + + if dist.is_initialized() and dist.get_world_size(comm_group) > 1: + ixfd.all_reduce(out, group=comm_group, async_op=True) + + return out diff --git a/ixformer_sdk/inference/overlap/linear_mlp_overlap_comm.py b/ixformer_sdk/inference/overlap/linear_mlp_overlap_comm.py new file mode 100644 index 00000000..e6bc8d67 --- /dev/null +++ b/ixformer_sdk/inference/overlap/linear_mlp_overlap_comm.py @@ -0,0 +1,305 @@ +from contextlib import nullcontext +from typing import List + +import torch.cuda + +from ...distributed import _distributed as ixfd +from ...distributed import overlap_comm as base_overlap_comm +from ...distributed.overlap_comm import GemmAllReduceSplitOverlapComm +from .. import overlap as overlap_base + + +class LinearMLPOverlapCommHook: + def on_mlp_linear2_finished( + self, + overlap_comm: "LinearMLPOverlapComm", + num_chunks, + chunk_idx, + hidden_states_chunk, + residual_chunk, + ): + pass + + def on_mlp_finished( + self, + overlap_comm: "LinearMLPOverlapComm", + hidden_states_chunks, + residual_chunks, + ): + pass + + +class LinearMLPOverlapComm(GemmAllReduceSplitOverlapComm): + def __init__(self, *args, **kwargs): + super().__init__(num_compute_streams=1, *args, **kwargs) + + self._mlp_linear1_start_events: List[torch.cuda.Event] = [ + torch.cuda.Event() for _ in range(self.num_chunks) + ] + self._mlp_linear1_end_events: List[torch.cuda.Event] = [ + torch.cuda.Event() for _ in range(self.num_chunks) + ] + + self._mlp_linear1_stream: torch.cuda.Stream = torch.cuda.Stream() + + def stop_linear_comm(self, chunk_idx): + event = self._mlp_linear1_start_events[chunk_idx] + event.record(self._comm_stream) + + def start_mlp_linear1(self, chunk_idx): + self._mlp_linear1_stream.wait_event(self._mlp_linear1_start_events[chunk_idx]) + + def stop_mlp_linear1(self, chunk_idx): + event = self._mlp_linear1_end_events[chunk_idx] + event.record(self._mlp_linear1_stream) + + def start_mlp_linear2(self, chunk_idx): + compute_stream = self._compute_streams[chunk_idx % self.num_compute_streams] + compute_stream.wait_event(self._mlp_linear1_end_events[chunk_idx]) + + def compute( + self, + protocol: "overlap_base.LlamaDecoderLayerOverlapDefault", + attn_output, + residual, + *, + mlp_linear2_finished_callback=None, + mlp_finished_callback=None, + ): + """ """ + attn_output_shape = attn_output.shape + residual_shape = None if residual is None else residual.shape + + is_update_shape = attn_output.ndim > 2 + batch = 1 + if attn_output.ndim == 2: + seqlen = attn_output_shape[0] + else: + batch = attn_output_shape[0] + seqlen = attn_output_shape[1] + + parallel_dims = batch * seqlen + + if is_update_shape: + attn_output = attn_output.reshape(parallel_dims, -1) + if residual is not None: + residual = residual.reshape(-1, residual_shape[-1]) + + attn_output_chunks, residual_chunks = protocol.split_mlp_inputs( + attn_output, residual, self.num_chunks + ) + + out = protocol.create_mlp_output() + out_chunks = protocol.split_mlp_output(out, self.num_chunks) + + res_chunks = [] + + # 1. output project linear + for chunk_idx, (attn_output_chunk, residual_chunk) in enumerate( + zip(attn_output_chunks, residual_chunks) + ): + with self.compute_stream_context(chunk_idx): + hidden_states = protocol.attn_output_proj_linear( + self.num_chunks, + chunk_idx, + attn_output_chunk, + use_limited_gemm=chunk_idx != 0, + ) + + self.start_comm(chunk_idx) + ixfd.all_reduce( + hidden_states, + async_op=True, + group=self.comm_group, + use_comm_stream=True, + ) + self.stop_linear_comm(chunk_idx) + attn_output_chunks[chunk_idx] = hidden_states + + # 2. ln, mlp_linear1 and act + for chunk_idx, (hidden_states, residual_chunk) in enumerate( + zip(attn_output_chunks, residual_chunks) + ): + self.start_mlp_linear1(chunk_idx) + with self.stream_context(self._mlp_linear1_stream): + ( + hidden_states, + residual_chunk, + ) = protocol.attn_output_proj_linear_layer_norm( + self.num_chunks, chunk_idx, hidden_states, residual_chunk + ) + + hidden_states = protocol.mlp_linear1( + self.num_chunks, chunk_idx, hidden_states, use_limited_gemm=True + ) + hidden_states = protocol.mlp_activation(hidden_states) + + attn_output_chunks[chunk_idx] = hidden_states + res_chunks.append(residual_chunk) + + self.stop_mlp_linear1(chunk_idx) + + # 3. mlp_linear2 + for chunk_idx, hidden_states in enumerate(attn_output_chunks): + self.start_mlp_linear2(chunk_idx) + with self.compute_stream_context(chunk_idx): + hidden_states = protocol.mlp_linear2( + self.num_chunks, + chunk_idx, + hidden_states, + out=out_chunks[chunk_idx], + use_limited_gemm=True, + ) + + self.start_comm(chunk_idx) + ixfd.all_reduce( + hidden_states, + async_op=True, + group=self.comm_group, + use_comm_stream=True, + ) + + if mlp_linear2_finished_callback is not None: + mlp_linear2_finished_callback( + self, + self.num_chunks, + chunk_idx, + hidden_states, + res_chunks[chunk_idx], + ) + + if mlp_finished_callback is not None: + mlp_finished_callback(self, out_chunks, res_chunks) + + if is_update_shape: + out = out.reshape(attn_output_shape) + if residual is not None: + residual = residual.reshape(residual_shape) + + return out, residual + + def gemm_dispatcher( + self, + chunk_idx, + chunk_input, + weight, + chunk_out, + use_limited_gemm=False, + user_gemm_method=None, + *args, + **kwargs, + ): + if user_gemm_method is not None and callable(user_gemm_method): + ctx = self.ixf_limited_gemm_ctx if use_limited_gemm else nullcontext() + with ctx: + return user_gemm_method( + chunk_input, weight, out=chunk_out, *args, **kwargs + ) + + ctx = self.limited_gemm_ctx if use_limited_gemm else nullcontext() + with ctx: + return torch.matmul(chunk_input, weight.T, out=chunk_out) + + @classmethod + def is_supported(cls, input, num_chunks, comm_group): + if not cls.enable(): + return False + + ndim = input.ndim + shape = input.shape + + if ndim == 1: + m, k = 1, shape[0] + elif ndim == 2: + m, k = shape + else: + m, k = sum(shape[:-1]), shape[-1] + + return m >= 512 + + @classmethod + def native_forward( + cls, + attn_output, + residual, + linear_weight, + ln_layer, + mlp_weight1, + mlp_weight2, + mlp_activation, + linear_method=None, + mlp_linear1_method=None, + mlp_linear2_method=None, + group=None, + *args, + **kwargs, + ): + import ixformer.functions as ixff + + linear_method = linear_method or ixff.linear + mlp_linear1_method = mlp_linear1_method or ixff.linear + mlp_linear2_method = mlp_linear2_method or ixff.linear + + hidden_states = linear_method(attn_output, linear_weight) + ixfd.all_reduce(hidden_states, async_op=True, group=group) + + if ln_layer is not None: + hidden_states, residual = ln_layer(hidden_states, residual) + + hidden_states = mlp_linear1_method(hidden_states, mlp_weight1) + hidden_states = mlp_activation(hidden_states) + hidden_states = mlp_linear2_method(hidden_states, mlp_weight2) + ixfd.all_reduce(hidden_states, async_op=True, group=group) + + return hidden_states, residual + + +_DEFAULT_OVERLAP_GROUP = None +_DEFAULT_OVERLAP_COMM_N2 = None +_DEFAULT_OVERLAP_COMM_N4 = None +_DEFAULT_OVERLAP_CHUNKS = base_overlap_comm._DEFAULT_OVERLAP_CHUNKS + + +def linear_mlp_overlap( + protocol: "overlap_base.LlamaDecoderLayerOverlapProtocol", + attn_output, + residual, + num_chunks=None, + group=None, + *, + mlp_linear2_finished_callback=None, + mlp_finished_callback=None, +): + num_chunks = num_chunks or _DEFAULT_OVERLAP_CHUNKS + + global _DEFAULT_OVERLAP_GROUP + global _DEFAULT_OVERLAP_COMM_N2 + global _DEFAULT_OVERLAP_COMM_N4 + + if _DEFAULT_OVERLAP_GROUP is None: + _DEFAULT_OVERLAP_GROUP = group + + if num_chunks == 2 and group == _DEFAULT_OVERLAP_GROUP: + if _DEFAULT_OVERLAP_COMM_N2 is None: + _DEFAULT_OVERLAP_COMM_N2 = LinearMLPOverlapComm.dispatcher( + num_chunks=num_chunks, comm_group=group + ) + overlap_comm = _DEFAULT_OVERLAP_COMM_N2 + elif num_chunks == 4 and group == _DEFAULT_OVERLAP_GROUP: + if _DEFAULT_OVERLAP_COMM_N4 is None: + _DEFAULT_OVERLAP_COMM_N4 = LinearMLPOverlapComm.dispatcher( + num_chunks=num_chunks, comm_group=group + ) + overlap_comm = _DEFAULT_OVERLAP_COMM_N4 + else: + overlap_comm = LinearMLPOverlapComm.dispatcher( + num_chunks=num_chunks, comm_group=group + ) + + return overlap_comm.forward( + protocol, + attn_output, + residual, + mlp_linear2_finished_callback=mlp_linear2_finished_callback, + mlp_finished_callback=mlp_finished_callback, + ) diff --git a/ixformer_sdk/inference/overlap/llama_decoder_layer_overlap.py b/ixformer_sdk/inference/overlap/llama_decoder_layer_overlap.py new file mode 100644 index 00000000..9b90e252 --- /dev/null +++ b/ixformer_sdk/inference/overlap/llama_decoder_layer_overlap.py @@ -0,0 +1,1396 @@ +import enum +import typing +from abc import abstractmethod +from typing import Any, Callable, List, Optional, Tuple + +import ixformer._C.infer as ops +import ixformer.functions as ixff +import torch +import torch.distributed as dist + +import ixformer.distributed as ixfd +from ixformer.core import config + +from ...distributed.overlap_comm import GemmWithLimitedBlock +from .linear_mlp_overlap_comm import LinearMLPOverlapComm +from .overlap_comm import DecoderLayerOverlapComm + +KVCache = Tuple[torch.Tensor, torch.Tensor] + +OptionalTensor = typing.Union[None, torch.Tensor] + + +class LlamaDecoderLayerParams: + def __init__( + self, + *, + pre_input_layer_norm_weight: OptionalTensor, + qkv_linear_weight: torch.Tensor, + qkv_linear_bias: OptionalTensor, + attn_output_proj_linear_weight: torch.Tensor, + attn_output_proj_linear_bias: OptionalTensor, + attn_output_proj_linear_layer_norm_weight: torch.Tensor, + mlp_linear1_weight: torch.Tensor, + mlp_linear1_bias: OptionalTensor, + mlp_linear2_weight: torch.Tensor, + mlp_linear2_bias: OptionalTensor, + mlp_activation: Callable[[torch.Tensor, OptionalTensor], Any], + layer_norm_eps: float = 1e-5, + quant_mode: Optional[str] = None, + **kwargs, + ): + self.pre_input_layer_norm_weight = pre_input_layer_norm_weight + + self.qkv_linear_weight = qkv_linear_weight + self.qkv_linear_bias = qkv_linear_bias + + self.attn_output_proj_linear_weight = attn_output_proj_linear_weight + self.attn_output_proj_linear_bias = attn_output_proj_linear_bias + self.attn_output_proj_linear_layer_norm_weight = ( + attn_output_proj_linear_layer_norm_weight + ) + + self.mlp_linear1_weight = mlp_linear1_weight + self.mlp_linear1_bias = mlp_linear1_bias + + self.mlp_linear2_weight = mlp_linear2_weight + self.mlp_linear2_bias = mlp_linear2_bias + + self.mlp_activation = mlp_activation + + self.layer_norm_eps = layer_norm_eps + + self.quant_mode = quant_mode + + for k, v in kwargs: + setattr(self, k, v) + + @classmethod + def create_from(cls, params: "LlamaDecoderLayerParams", **extra_params): + param_attrs = dict(**params.__dict__) + param_attrs.update(**extra_params) + + return cls(**param_attrs) + + +class LlamaDecoderLayerOverlapProtocol: + GLOBAL_ENABLE_OVERLAP_CACHE = False + + class HookStage(enum.IntEnum): + kExited = 0 + kTracing = 1 + + class HookState: + def __init__(self, max_num_chunks): + self.max_num_chunks = max_num_chunks + self.stage = DecoderLayerOverlapComm.HookStage.kExited + + self.mlp_linaer2_end_events = [ + torch.cuda.Event() for _ in range(max_num_chunks) + ] + self.ln_attn_end_event = torch.cuda.Event() + + self.overlap_comm: Optional[LinearMLPOverlapComm] = None + + def is_tracing_stage(self): + return self.stage == DecoderLayerOverlapComm.HookStage.kTracing + + def enter(self, overlap_comm, chunk_idx): + self.stage = DecoderLayerOverlapComm.HookStage.kTracing + + self.overlap_comm = overlap_comm + self.mlp_linaer2_end_events[chunk_idx].record(overlap_comm._comm_stream) + + def exit(self): + self.overlap_comm = None + self.stage = DecoderLayerOverlapComm.HookStage.kExited + + def __str__(self): + return f"HookState(overlap_comm={self.overlap_comm}, stage={self.stage})" + + def __repr__(self): + return self.__str__() + + _overlap_comm_hook_state = dict() + + def __init__( + self, + model_id: int, + layer_idx: int, + attn_q_size: int, + attn_kv_size: int, + params: LlamaDecoderLayerParams = None, + max_num_chunks=4, + ): + """ + DecoderLayer 的流程： + ln_qkv: InputLayerNorm(hidden_states, [residual]) -> qkv_proj(hidden_states) -> q, k, v = split(hidden_states) -> Attention(q, k, v) + linear_mlp: AttentionOutputProj(hidden_states) -> PostLayerNorm(hidden_states) -> MLPLinear1 -> MLPActivation -> MLPLinear2 + + 其中：AttentionOutputProj 和 MLPLinear2 之后如果使用 TP，那么需要进行 AllReduce + + 通过上述流程，该类的目的是将 MLPLinear2 后的 AllReduce 和 DecoderLayer 最开始的 ln_qkv 进行 Overlap。 + 其中，第一层 DecoderLayer 不进行 ln_qkv 的 Overlap，因为在第一层之前没有通讯。 + 我们需要将第 i 层 MLPLinear2 后的通讯 和 第 i + 1 层的 ln_qkv 进行 Overlap。 + + 为了管理当前的状态和获取前一层的状态，从而设计了 DecoderLayerOverlapComm 类。 + 该类需要 model_id 来推断当前正在运行的模型，用 layer_idx 来标记每一层的开始和结束， + 以及通过 layer_idx 去获取前一层的状态。 + + 注： + - 在 call_ln_qkv_overlap 中对 Tensor 进行切分时， + 需要保持和 linear_mlp 切分的大小是一致的，否则会出现 Tensor 的数据不对应； + - 如果需要使用 ln_qkv 进行 Overlap，那么必须使用该类的 linear_mlp 去替换 linear_mlp_overlap + + :param model_id: 模型的 id，可以使用 id(model) 去设置 + :param layer_idx: layer 的索引，注意，需要从 0 到 NumLayers 的顺序去完成构造 + :param max_num_chunks: 最大能进行切分的次数 + """ + + if layer_idx < 0: + raise RuntimeError(f"Invalid layer_idx, got {layer_idx}.") + self._model_id = model_id + self._layer_idx = layer_idx + self._max_num_chunks = max_num_chunks + + self._state = self.HookState(max_num_chunks) + self._overlap_comm_hook_state[(model_id, layer_idx)] = self._state + self._prev_layer_state = ( + None + if layer_idx == 0 + else self._overlap_comm_hook_state[(model_id, layer_idx - 1)] + ) + + self.params: LlamaDecoderLayerParams = params + + self.attn_q_size = attn_q_size + self.attn_kv_size = attn_kv_size + + self.limited_blas_gemm_ctx = GemmWithLimitedBlock() + + self._current_comm_group = None + + def set_params(self, params: LlamaDecoderLayerParams): + self.params = params + + @property + def model_id(self): + return self._model_id + + @property + def layer_idx(self): + return self._layer_idx + + @property + def max_num_chunks(self): + return self._max_num_chunks + + @property + def state(self) -> "DecoderLayerOverlapComm.HookState": + return self._state + + @property + def prev_layer_state(self) -> "DecoderLayerOverlapComm.HookState": + return self._prev_layer_state + + @property + def out_last_dim(self): + return self.attn_q_size + 2 * self.attn_kv_size + + @classmethod + def is_supported(cls, input, num_chunks, comm_group): + return ( + config.IXFORMER_ENABLE_OVERLAP_COMM + and LinearMLPOverlapComm.is_supported(input, num_chunks, comm_group) + ) + + def forward( + self, + self_attn: Callable, + positions: torch.Tensor, + hidden_states: torch.Tensor, + *, + num_chunks=None, + group=None, + residual: Optional[torch.Tensor], + self_attn_kwargs: dict, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + 等价于 DecoderLayer.forward + + :param self_attn: Function(Q, K, V, **self_attn_kwargs), 计算 SelfAttention 的输出， + Q, K, V 会在 self.attention 中根据 q_size 和 kv_size 进行分割得到， + 如果需要额外的参数，可以使用 self_attn_kwargs 进行传递. + :param positions: 如果使用 RotaryEmbedding，那么会被传入到该函数中 + :param hidden_states: [Batch * SeqLen, HiddenSize], 前一层的输出 + :param num_chunks: 在 Overlap 时分块数量 + :param group: 通讯组 + :param residual: 前一层的残差 + :param self_attn_kwargs: self_attn 函数的额外参数 + :return: DecoderLayerOut[Batch * SeqLen, HiddenSize], Residual[Batch * SeqLen, HiddenSize] + """ + + self._current_comm_group = group + + # 仅在第一层 Layer 去判断是否使用 Overlap， + # 如果第一层 Layer 启用，那么后面的所有 Layer 也都会使用 Overlap + if self.layer_idx == 0: + self.__class__.GLOBAL_ENABLE_OVERLAP_CACHE = self.is_supported( + hidden_states, num_chunks, comm_group=group + ) + + enable_overlap = self.__class__.GLOBAL_ENABLE_OVERLAP_CACHE + + qkv, residual = self.ln_qkv( + hidden_states, residual, enable_overlap=enable_overlap + ) + + attn_output = self.attention( + qkv, residual, self_attn=self_attn, positions=positions, **self_attn_kwargs + ) + + attn_output, residual = self.linear_mlp( + attn_output, + residual, + num_chunks=num_chunks, + group=group, + enable_overlap=enable_overlap, + ) + + return attn_output, residual + + def is_ln_qkv_overlap(self): + return not ( + self.layer_idx == 0 + or not self.prev_layer_state.is_tracing_stage() + or self.prev_layer_state.overlap_comm is None + ) + + def ln_qkv( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + enable_overlap: bool = False, + ) -> Tuple[torch.Tensor, OptionalTensor]: + """ + LayerNorm -> QKVLinear + + :param hidden_states: shape[Batch * SeqLen, HiddenSize] + :param residual: shape[Batch * SeqLen, HiddenSize] + :return: qkv, residual + """ + if self.is_ln_qkv_overlap() and enable_overlap: + qkv, residual = self.call_ln_qkv_overlap(hidden_states, residual) + else: + qkv, residual = self.call_ln_qkv(hidden_states, residual) + + return qkv, residual + + def call_ln_qkv_overlap(self, hidden_states, residual): + if hidden_states.ndim != 2: + raise RuntimeError( + f"Expected 2-dim for hidden state, but got {hidden_states.ndim}." + ) + + num_chunks = self.prev_layer_state.overlap_comm.num_chunks + overlap_comm: LinearMLPOverlapComm = self.prev_layer_state.overlap_comm + + if num_chunks > self.max_num_chunks: + raise RuntimeError( + f"The layer is not support more than {self.max_num_chunks}, got {num_chunks}." + ) + + hidden_state_chunks = list(torch.chunk(hidden_states, num_chunks, dim=0)) + if residual is None: + residual = hidden_states + residual_chunks = [None] * num_chunks + else: + residual_chunks = torch.chunk(residual, num_chunks, dim=0) + + out = torch.empty( + (hidden_states.shape[0], self.out_last_dim), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + out_chunks = list(torch.chunk(out, num_chunks, dim=0)) + + for chunk_idx, (hidden_state_chunk, residual_chunk, out_chunk) in enumerate( + zip(hidden_state_chunks, residual_chunks, out_chunks) + ): + overlap_comm._compute_streams[ + chunk_idx % overlap_comm.num_compute_streams + ].wait_event(self.prev_layer_state.mlp_linaer2_end_events[chunk_idx]) + with overlap_comm.compute_stream_context(chunk_idx): + self.call_ln_qkv( + hidden_state_chunk, + residual_chunk, + num_chunks, + chunk_idx, + use_limited_gemm=chunk_idx != (num_chunks - 1), + out=out_chunk, + ) + + self.prev_layer_state.exit() + overlap_comm.stop_overlap() + + out, residual = self.ln_qkv_linear_callback(out, residual) + + return out, residual + + def call_ln_qkv( + self, + hidden_state, + residual, + num_chunks=1, + chunk_idx=0, + use_limited_gemm=False, + out=None, + ): + pre_hidden_state = hidden_state + hidden_state, residual = self.pre_input_layer_norm( + num_chunks, chunk_idx, hidden_state, residual + ) + if residual is None: + residual = pre_hidden_state + + qkv = self.qkv_linear( + num_chunks, + chunk_idx, + hidden_state, + out=out, + use_limited_gemm=use_limited_gemm, + ) + + return qkv, residual + + def linear_mlp( + self, + attn_output: torch.Tensor, + residual: torch.Tensor, + num_chunks: int, + group=None, + enable_overlap: bool = False, + ) -> Tuple[torch.Tensor, OptionalTensor]: + """ + OProj -> AllReduce -> PostLayerNorm -> Linear1 -> Activation -> Linear2 -> AllReduce + + :param attn_output: Attention 的输出 + :param residual: 残差 + :param num_chunks: 在 Overlap 时，需要分为多少块 + :param group: 通讯组 + :param enable_overlap: 是否启用 overlap + :return: DecoderLayerOut[Batch * SeqLen, HiddenSize], Residual[Batch * SeqLen, HiddenSize] + """ + + if enable_overlap: + return ixff.linear_mlp_overlap( + self, + attn_output, + residual, + num_chunks=num_chunks, + group=group, + mlp_linear2_finished_callback=self.on_mlp_linear2_finished, + ) + + hidden_states = self.attn_output_proj_linear( + num_chunks=1, + chunk_idx=0, + attn_out_chunk=attn_output, + use_limited_gemm=False, + ) + ixfd.all_reduce(hidden_states, async_op=True, group=self._current_comm_group) + + hidden_states, residual = self.attn_output_proj_linear_layer_norm( + num_chunks=1, chunk_idx=0, hidden_states=hidden_states, residual=residual + ) + + hidden_states = self.mlp_linear1( + num_chunks=1, + chunk_idx=0, + hidden_states=hidden_states, + use_limited_gemm=False, + ) + + hidden_states = self.mlp_activation(hidden_states) + + hidden_states = self.mlp_linear2( + num_chunks=1, + chunk_idx=0, + hidden_states=hidden_states, + out=None, + use_limited_gemm=False, + ) + ixfd.all_reduce(hidden_states, async_op=True, group=self._current_comm_group) + + return hidden_states, residual + + def on_mlp_linear2_finished( + self, + overlap_comm: LinearMLPOverlapComm, + num_chunks, + chunk_idx, + hidden_states_chunk, + residual_chunk, + ): + """ + 这是一个 LinearMLPOverlapComm 的回调函数，在 MLP Linear2 后面的通讯结束时被执行， + 在这里是为了将 MLP Linear2 后面的通讯和 LnQKV 进行 Overlap，需要进行 cuda event 的同步。 + """ + self.state.enter(overlap_comm, chunk_idx) + + def _gemm_dispatcher( + self, + input: torch.Tensor, + weight: torch.Tensor, + bias: OptionalTensor = None, + input_scales: OptionalTensor = None, + smooth_scales: OptionalTensor = None, + weight_scales: OptionalTensor = None, + out: OptionalTensor = None, + use_limited_gemm: bool = False, + quant_group_size: int = -1, + out_dtype: Optional[torch.dtype] = None, + ): + """ + 调用不同精度的 gemm，已验证 fp16，w8a8 + """ + + # 下面判断的顺序不能随意改变 + + # float + if input_scales is None and weight_scales is None: + if out is not None and out.dtype not in [ + torch.half, + torch.bfloat16, + torch.float, + ]: + raise RuntimeError( + f"linear is supported half or float, but got {out.dtype}." + ) + + # if use_limited_gemm: + # with self.limited_blas_gemm_ctx: + # out = ops.cublas_linear( + # input, weight, bias=bias, out=out, persistent=use_limited_gemm + # ) + # else: + out = ixff.linear(input, weight, bias=bias, output=out, persistent=use_limited_gemm) + return out + + elif weight_scales is None: + raise RuntimeError(f"got invalid quantized weight scales, got none.") + + # smmoth quant with w8a8 + elif (input_scales is None and smooth_scales is not None) or self.params.quant_mode in ["smoothquant", "compressed_tensors"]: + x_shape = input.shape + dtype = input.dtype + if self.params.quant_mode == "compressed_tensors": + x, x_scales = ixff.scaled_int8_quant(input, smooth_scales) + else: + x, x_scales = ixff.dynamic_scaled_quant_dynamic_int8(input, smooth_scales) + + x = ixff.w8a8( + input=x, + weight=weight, + i_scales=x_scales, + w_scales=weight_scales, + output=out, + persistent=use_limited_gemm, + out_dtype=dtype, + ) + out = x.view(*x_shape[:-1], -1) + + # w8a16 + elif input_scales is None and weight_scales is not None: + if out is not None and out.dtype not in [ + torch.half, + torch.bfloat16, + torch.float, + ]: + raise RuntimeError( + f"w8a16 is supported half or float, but got {out.dtype}." + ) + + out = ixff.w8a16( + input, weight, weight_scales, output=out, group_size=quant_group_size, persistent=int(use_limited_gemm) + ) + + # w8a8 + elif input_scales is not None and weight_scales is not None: + if out is not None and out.dtype not in [ + torch.half, + torch.bfloat16, + torch.float, + ]: + raise RuntimeError( + f"w8a8 is supported half or float, but got {out.dtype}." + ) + + out = ixff.w8a8( + input, + weight, + input_scales, + weight_scales, + output=out, + persistent=use_limited_gemm, + out_dtype=out_dtype, + ) + + else: + raise RuntimeError("dispatcher gemm fail.") + + if bias is None: + return out + return out + bias + + @abstractmethod + def split_ln_qkv_input( + self, hidden_states: torch.Tensor, residual: OptionalTensor, num_chunks: int + ) -> Tuple[List[torch.Tensor], List[OptionalTensor]]: + """ + 对 ln_qkv 的输入进行切分，仅被用在 overlap 时 + :param hidden_states: 对 hidden_states 进行切分 + :param residual: 对 residual 进行切分，如果 residual 为 None，那么应该返回 [None] * num_chunks + :param num_chunks: 分块数量 + :return: HiddenStatesChunks, ResidualChunks + """ + raise NotImplementedError() + + @abstractmethod + def create_ln_qkv_output(self, num_chunks: int) -> torch.Tensor: + """ + 创建 ln_qkv 的输出，仅被用在 overlap 时 + :param num_chunks: 分块数量 + :return: Tensor + """ + raise NotImplementedError() + + @abstractmethod + def split_ln_qkv_output( + self, out: torch.Tensor, num_chunks: int + ) -> List[torch.Tensor]: + """ + 对上面创建的输出 Tensor 进行切分，仅被用在 overlap 时 + :param out: 上面创建的 Tensor + :param num_chunks: 分块数量 + :return: OutChunks + """ + raise NotImplementedError() + + @abstractmethod + def pre_input_layer_norm( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + residual: OptionalTensor, + ) -> Tuple[torch.Tensor, OptionalTensor]: + """ + DecoderLayer 中的第一个 LayerNorm + :param num_chunks: 分块数量 + :param chunk_idx: 分块的索引 + :param hidden_states: 分块后的输入 + :param residual: 分块后的残差 + :return: HiddenStates，Residual + """ + raise NotImplementedError() + + @abstractmethod + def qkv_linear( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + out: OptionalTensor, + use_limited_gemm=False, + ) -> torch.Tensor: + """ + DecoderLayer 中的 QkvLinear + + :param num_chunks: 分块数量 + :param chunk_idx: 分块的索引 + :param hidden_states: linear 的输入 + :param out: linear 的输出 + :param use_limited_gemm: 是否限制 gemm 的计算资源 + :return: Out + """ + raise NotImplementedError() + + @abstractmethod + def ln_qkv_linear_callback( + self, qkv: torch.Tensor, residual: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + return qkv, residual + + @abstractmethod + def attention( + self, + qkv: torch.Tensor, + residual: torch.Tensor, + *, + self_attn: Callable, + positions: OptionalTensor, + **kwargs, + ) -> torch.Tensor: + """ + 计算 Attention + :param qkv: QkvLinear 的输出 + :param residual: PreLayerNorm 输出的残差 + :param self_attn: 计算 SelfAttention 的函数 + :param positions: 位置编码，被用在 RotaryEmbedding 中 + :param kwargs: self_attn 的额外参数 + :return: Attention 的输出 + """ + raise NotImplementedError() + + @abstractmethod + def create_mlp_output(self) -> torch.Tensor: + """ + 创建 MLP 的输出，仅被用在 overlap 时 + """ + + raise NotImplementedError() + + @abstractmethod + def split_mlp_output( + self, out: torch.Tensor, num_chunks: int + ) -> List[torch.Tensor]: + """ + 对上面创建的输出进行分块，仅被用在 overlap 时 + :param out: 上面函数的输出 + :param num_chunks: 分块数量 + :return: OutChunks + """ + raise NotImplementedError() + + @abstractmethod + def split_mlp_inputs( + self, attn_out: torch.Tensor, residual: torch.Tensor, num_chunks: int + ) -> Tuple[List[torch.Tensor], List[OptionalTensor]]: + """ + 对 MLP 的输入进行分块，仅被用在 overlap 时 + :param attn_out: Attention 的输出 + :param residual: 残差 + :param num_chunks: 分块数量 + :return: AttnOutChunks, ResidualChunks + """ + raise NotImplementedError() + + @abstractmethod + def attn_output_proj_linear( + self, + num_chunks: int, + chunk_idx: int, + attn_out_chunk: torch.Tensor, + use_limited_gemm: bool = False, + ) -> torch.Tensor: + """ + Attention 后面的 o_proj Linear + :param num_chunks: 分块数量 + :param chunk_idx: 分块的索引 + :param attn_out_chunk: Linear 的输入 + :param use_limited_gemm: 是否限制 gemm 的计算资源 + :return: Out + """ + raise NotImplementedError() + + @abstractmethod + def attn_output_proj_linear_layer_norm( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ) -> Tuple[torch.Tensor, OptionalTensor]: + """ + attention output projection linear 后面的 LayerNorm + :param num_chunks: 分块数量 + :param chunk_idx: 分块索引 + :param hidden_states: LN 的输入 + :param residual: 残差 + :return: HiddenStates，Residual + """ + raise NotImplementedError() + + @abstractmethod + def mlp_linear1( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + use_limited_gemm: bool = False, + ) -> torch.Tensor: + """ + MLP 中的第一个 Linear + :param num_chunks: 分块数量 + :param chunk_idx: 分块索引 + :param hidden_states: Linear 的输入 + :param use_limited_gemm: 是否限制 gemm 的计算资源 + :return: Out + """ + raise NotImplementedError() + + @abstractmethod + def mlp_activation(self, hidden_states: torch.Tensor) -> torch.Tensor: + """ + MLP 中的激活函数 + """ + raise NotImplementedError() + + @abstractmethod + def mlp_linear2( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + out: OptionalTensor = None, + use_limited_gemm: bool = False, + ) -> torch.Tensor: + """ + MLP 中的第二个 Linear + :param num_chunks: 分块数量 + :param chunk_idx: 分块索引 + :param hidden_states: Linear 的输入 + :param out: Linear 的输出 + :param use_limited_gemm: 是否限制 gemm 的计算资源 + :return: Out + """ + raise NotImplementedError() + + @abstractmethod + def mlp_callback(self, mlp_out: torch.Tensor, residual: OptionalTensor): + return mlp_out, residual + + +class LlamaDecoderLayerOverlapDefault(LlamaDecoderLayerOverlapProtocol): + def split_ln_qkv_input( + self, hidden_states: torch.Tensor, residual: OptionalTensor, num_chunks: int + ) -> Tuple[List[torch.Tensor], List[OptionalTensor]]: + self.input_hidden_states_shape = list(hidden_states.shape) + self.input_hidden_states_device = hidden_states.device + self.input_hidden_states_dtype = hidden_states.dtype + + hidden_states_chunks = list(torch.chunk(hidden_states, num_chunks, dim=0)) + if residual is None: + residual_chunks = [None] * num_chunks + else: + residual_chunks = list(torch.chunk(residual, num_chunks, dim=0)) + + return hidden_states_chunks, residual_chunks + + def create_ln_qkv_output(self, num_chunks: int) -> torch.Tensor: + return torch.empty( + (self.input_hidden_states_shape[0], self.out_last_dim), + device=self.input_hidden_states_device, + dtype=self.input_hidden_states_dtype, + ) + + def split_ln_qkv_output( + self, out: torch.Tensor, num_chunks: int + ) -> List[torch.Tensor]: + return list(torch.chunk(out, num_chunks, dim=0)) + + def pre_input_layer_norm( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + residual: OptionalTensor, + ) -> Tuple[torch.Tensor, OptionalTensor]: + if residual is None: + return ( + ixff.rms_norm( + hidden_states, + self.params.pre_input_layer_norm_weight, + eps=self.params.layer_norm_eps, + ), + None, + ) + else: + ixff.residual_rms_norm( + hidden_states, + residual, + self.params.pre_input_layer_norm_weight, + eps=self.params.layer_norm_eps, + ) + return hidden_states, residual + + def qkv_linear( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + out: OptionalTensor, + use_limited_gemm=False, + ) -> torch.Tensor: + return self._gemm_dispatcher( + input=hidden_states, + weight=self.params.qkv_linear_weight, + bias=self.params.qkv_linear_bias, + out=out, + use_limited_gemm=use_limited_gemm, + ) + + def attention( + self, + qkv: torch.Tensor, + residual: torch.Tensor, + *, + self_attn: Callable, + positions: OptionalTensor, + rotary_embedding: Callable = None, + **kwargs, + ): + q, k, v = qkv.split( + [self.attn_q_size, self.attn_kv_size, self.attn_kv_size], + dim=-1, + ) + if rotary_embedding is not None: + q, k = rotary_embedding(positions, q, k) + attn_output = self_attn(q, k, v, **kwargs) + + self.attn_output_shape = attn_output.shape + self.attn_output_device = attn_output.device + self.attn_output_dtype = attn_output.dtype + + return attn_output + + def create_mlp_output(self) -> torch.Tensor: + return torch.empty( + [self.attn_output_shape[0], self.params.mlp_linear2_weight.shape[0]], + device=self.attn_output_device, + dtype=self.attn_output_dtype, + ) + + def split_mlp_output( + self, out: torch.Tensor, num_chunks: int + ) -> List[torch.Tensor]: + return list(torch.chunk(out, num_chunks, dim=0)) + + def split_mlp_inputs( + self, attn_out: torch.Tensor, residual: torch.Tensor, num_chunks: int + ) -> Tuple[List[torch.Tensor], List[OptionalTensor]]: + attn_output_chunks = list(torch.chunk(attn_out, num_chunks, dim=0)) + if residual is None: + residual_chunks = [None] * num_chunks + else: + residual_chunks = torch.chunk(residual, num_chunks, dim=0) + + return attn_output_chunks, residual_chunks + + def attn_output_proj_linear( + self, + num_chunks: int, + chunk_idx: int, + attn_out_chunk: torch.Tensor, + use_limited_gemm: bool = False, + ) -> torch.Tensor: + return self._gemm_dispatcher( + input=attn_out_chunk, + weight=self.params.attn_output_proj_linear_weight, + bias=self.params.attn_output_proj_linear_bias, + use_limited_gemm=use_limited_gemm, + ) + + def attn_output_proj_linear_layer_norm( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ) -> Tuple[torch.Tensor, OptionalTensor]: + if residual is None: + return ( + ixff.rms_norm( + hidden_states, + self.params.attn_output_proj_linear_layer_norm_weight, + eps=self.params.layer_norm_eps, + ), + None, + ) + else: + return ixff.residual_rms_norm( + hidden_states, + residual, + self.params.attn_output_proj_linear_layer_norm_weight, + self.params.layer_norm_eps, + ) + + @abstractmethod + def mlp_linear1( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + use_limited_gemm: bool = False, + ) -> torch.Tensor: + return self._gemm_dispatcher( + input=hidden_states, + weight=self.params.mlp_linear1_weight, + bias=self.params.mlp_linear1_bias, + use_limited_gemm=use_limited_gemm, + ) + + def mlp_activation(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.params.mlp_activation(hidden_states) + + def mlp_linear2( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + out: OptionalTensor = None, + use_limited_gemm: bool = False, + ) -> torch.Tensor: + return self._gemm_dispatcher( + input=hidden_states, + weight=self.params.mlp_linear2_weight, + bias=self.params.mlp_linear2_bias, + out=out, + use_limited_gemm=use_limited_gemm, + ) + + +class LlamaDecoderLayerParamsQuant(LlamaDecoderLayerParams): + def __init__( + self, + *, + # qkv + qkv_linear_smooth_scales: OptionalTensor, + qkv_linear_weight_scales: OptionalTensor, + qkv_linear_quant_group_size: Optional[int] = -1, + # attn_output_proj + attn_output_proj_linear_smooth_scales: OptionalTensor, + attn_output_proj_linear_weight_scales: OptionalTensor, + attn_output_proj_linear_quant_group_size: Optional[int] = -1, + # mlp_linear1 + mlp_linear1_smooth_scales: OptionalTensor, + mlp_linear1_weight_scales: OptionalTensor, + mlp_linear1_quant_group_size: Optional[int] = -1, + # mlp_linear2 + mlp_linear2_smooth_scales: OptionalTensor, + mlp_linear2_weight_scales: OptionalTensor, + mlp_linear2_quant_group_size: Optional[int] = -1, + # other + activation_dtype: Optional[torch.dtype] = None, + **kwargs, + ): + super().__init__(**kwargs) + + self.qkv_linear_smooth_scales = qkv_linear_smooth_scales + self.qkv_linear_weight_scales = qkv_linear_weight_scales + self.qkv_linear_quant_group_size = qkv_linear_quant_group_size + + self.attn_output_proj_linear_smooth_scales = ( + attn_output_proj_linear_smooth_scales + ) + self.attn_output_proj_linear_weight_scales = ( + attn_output_proj_linear_weight_scales + ) + self.attn_output_proj_linear_quant_group_size = ( + attn_output_proj_linear_quant_group_size + ) + + self.mlp_linear1_smooth_scales = mlp_linear1_smooth_scales + self.mlp_linear1_weight_scales = mlp_linear1_weight_scales + self.mlp_linear1_quant_group_size = mlp_linear1_quant_group_size + + self.mlp_linear2_smooth_scales = mlp_linear2_smooth_scales + self.mlp_linear2_weight_scales = mlp_linear2_weight_scales + self.mlp_linear2_quant_group_size = mlp_linear2_quant_group_size + + self.activation_dtype = activation_dtype + + +class LlamaDecoderLayerOverlapQuant(LlamaDecoderLayerOverlapDefault): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.qkv_linear_input_scales = [None] * self.max_num_chunks + self.attn_output_proj_linear_input_scales = [None] * self.max_num_chunks + self.mlp_linear1_input_scales = [None] * self.max_num_chunks + self.mlp_linear2_input_scales = [None] * self.max_num_chunks + + def _dispatch_layer_norm( + self, + input: torch.Tensor, + weight: torch.Tensor, + residual: OptionalTensor, + smooth_scales: OptionalTensor, + ) -> typing.Union[ + Tuple[torch.Tensor, OptionalTensor], + Tuple[torch.Tensor, OptionalTensor, torch.Tensor], + ]: + if smooth_scales is None: + if residual is None: + return ( + ixff.rms_norm(input, weight, eps=self.params.layer_norm_eps), + None, + ) + else: + ixff.residual_rms_norm( + input, residual, weight, eps=self.params.layer_norm_eps + ) + return input, residual + + elif smooth_scales is not None: + if residual is None: + hidden_states, scales = ixff.residual_rms_norm_dynamic_int8( + input=input, + weight=weight, + smooth_scales=smooth_scales, + eps=self.params.layer_norm_eps, + ) + return hidden_states, None, scales + else: + hidden_states, residual, scales = ixff.residual_rms_norm_dynamic_int8( + input=input, + residual=residual, + weight=weight, + smooth_scales=smooth_scales, + eps=self.params.layer_norm_eps, + ) + return hidden_states, residual, scales + else: + raise RuntimeError("dispatcher layer norm fail.") + + def pre_input_layer_norm( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + residual: OptionalTensor, + ) -> Tuple[torch.Tensor, OptionalTensor]: + self.params: LlamaDecoderLayerParamsQuant + outs = self._dispatch_layer_norm( + hidden_states, + self.params.pre_input_layer_norm_weight, + residual, + self.params.qkv_linear_smooth_scales, + ) + + self.qkv_linear_input_scales[chunk_idx] = None + if len(outs) == 3: + self.qkv_linear_input_scales[chunk_idx] = outs[2] + + return outs[:2] + + def qkv_linear( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + out: OptionalTensor, + use_limited_gemm=False, + ) -> torch.Tensor: + self.params: LlamaDecoderLayerParamsQuant + return self._gemm_dispatcher( + input=hidden_states, + weight=self.params.qkv_linear_weight, + input_scales=self.qkv_linear_input_scales[chunk_idx], + weight_scales=self.params.qkv_linear_weight_scales, + out=out, + use_limited_gemm=use_limited_gemm, + quant_group_size=self.params.qkv_linear_quant_group_size, + out_dtype=self.params.activation_dtype, + ) + + def attn_output_proj_linear( + self, + num_chunks: int, + chunk_idx: int, + attn_out_chunk: torch.Tensor, + use_limited_gemm: bool = False, + ) -> torch.Tensor: + self.params: LlamaDecoderLayerParamsQuant + return self._gemm_dispatcher( + input=attn_out_chunk, + weight=self.params.attn_output_proj_linear_weight, + input_scales=self.attn_output_proj_linear_input_scales[chunk_idx], + smooth_scales=self.params.attn_output_proj_linear_smooth_scales, + weight_scales=self.params.attn_output_proj_linear_weight_scales, + out=None, + use_limited_gemm=use_limited_gemm, + quant_group_size=self.params.attn_output_proj_linear_quant_group_size, + out_dtype=self.params.activation_dtype, + ) + + def attn_output_proj_linear_layer_norm( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ) -> Tuple[torch.Tensor, OptionalTensor]: + self.params: LlamaDecoderLayerParamsQuant + outs = self._dispatch_layer_norm( + hidden_states, + self.params.attn_output_proj_linear_layer_norm_weight, + residual, + self.params.mlp_linear1_smooth_scales, + ) + + self.mlp_linear1_input_scales[chunk_idx] = None + if len(outs) == 3: + self.mlp_linear1_input_scales[chunk_idx] = outs[2] + + return outs[:2] + + def mlp_linear1( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + use_limited_gemm: bool = False, + ) -> torch.Tensor: + self.params: LlamaDecoderLayerParamsQuant + self.mlp_linear1_num_chuns = num_chunks + self.mlp_linear1_chunk_idx = chunk_idx + return self._gemm_dispatcher( + input=hidden_states, + weight=self.params.mlp_linear1_weight, + input_scales=self.mlp_linear1_input_scales[chunk_idx], + weight_scales=self.params.mlp_linear1_weight_scales, + out=None, + use_limited_gemm=use_limited_gemm, + quant_group_size=self.params.mlp_linear1_quant_group_size, + out_dtype=self.params.activation_dtype, + ) + + def mlp_activation(self, hidden_states: torch.Tensor) -> torch.Tensor: + self.params: LlamaDecoderLayerParamsQuant + if self.params.mlp_linear2_smooth_scales is None: + self.mlp_linear2_input_scales[self.mlp_linear1_chunk_idx] = None + return self.params.mlp_activation(hidden_states) + + out, scales = self.params.mlp_activation( + hidden_states, self.params.mlp_linear2_smooth_scales + ) + self.mlp_linear2_input_scales[self.mlp_linear1_chunk_idx] = scales + return out + + def mlp_linear2( + self, + num_chunks: int, + chunk_idx: int, + hidden_states: torch.Tensor, + out: OptionalTensor = None, + use_limited_gemm: bool = False, + ) -> torch.Tensor: + self.params: LlamaDecoderLayerParamsQuant + return self._gemm_dispatcher( + input=hidden_states, + weight=self.params.mlp_linear2_weight, + input_scales=self.mlp_linear2_input_scales[chunk_idx], + weight_scales=self.params.mlp_linear2_weight_scales, + out=out, + use_limited_gemm=use_limited_gemm, + quant_group_size=self.params.mlp_linear2_quant_group_size, + out_dtype=self.params.activation_dtype, + ) + + +def is_vllm_supported_quant_mode(quant_mode: Optional[str]): + return quant_mode in [None, "smoothquant", "compressed_tensors"] + + +def get_vllm_llama_decoder_layer_protocol_cls( + layer: torch.nn.Module, quant_mode: Optional[str] +): + if quant_mode is None: + return LlamaDecoderLayerOverlapDefault + + elif quant_mode in ["smoothquant", "compressed_tensors"]: + return LlamaDecoderLayerOverlapQuant + + raise RuntimeError(f"got unsupported quantized mode: {quant_mode}.") + + +def create_vllm_llama_decoder_layer_params( + layer: torch.nn.Module, + quant_mode: Optional[str], + activation_dtype: Optional[torch.dtype] = None, +) -> LlamaDecoderLayerParams: + + # 在 vllm 的 w8a8 中，input 的排布为 [M, K], weight 的排布为 [K, N] (通过是否是 contiguous 来判断) + # 目前实现的 w8a8 不是该格式，需要将 weight 转置为 [N, K] + def transpose_weight(weight: torch.Tensor): + if not weight.is_contiguous() and weight.ndim == 2: + return weight.transpose(0, 1) + return weight + + T = transpose_weight + + def create_params_default(): + return LlamaDecoderLayerParams( + pre_input_layer_norm_weight=layer.input_layernorm.weight, + # qkv + qkv_linear_weight=T(layer.self_attn.qkv_proj.weight), + qkv_linear_bias=getattr(layer.self_attn.qkv_proj, "bias", None), + # attn_output_proj + attn_output_proj_linear_weight=T(layer.self_attn.o_proj.weight), + attn_output_proj_linear_bias=layer.self_attn.o_proj.bias, + # post layer norm + attn_output_proj_linear_layer_norm_weight=layer.post_attention_layernorm.weight, + # mlp linear1 + mlp_linear1_weight=T(layer.mlp.gate_up_proj.weight), + mlp_linear1_bias=getattr(layer.mlp.gate_up_proj, "bias", None), + # mlp linear2 + mlp_linear2_weight=T(layer.mlp.down_proj.weight), + mlp_linear2_bias=getattr(layer.mlp.down_proj, "bias", None), + mlp_activation=layer.mlp.act_fn, + layer_norm_eps=layer.input_layernorm.variance_epsilon, + ) + + if quant_mode is None: + return create_params_default() + + elif quant_mode in ["smoothquant", "compressed_tensors"]: + if activation_dtype is None: + raise RuntimeError( + "The smooth quantization need activation dtype as the output of gemm_w8a8." + ) + + weight_scales_key = "weight_scales" if quant_mode == "smoothquant" else "weight_scale" + smooth_scales_key = "smooth_scales" if quant_mode == "smoothquant" else "input_scale" + + model_params = create_params_default() + params = LlamaDecoderLayerParamsQuant.create_from( + model_params, + # qkv + qkv_linear_smooth_scales=getattr( + layer.self_attn.qkv_proj, smooth_scales_key, None + ), + qkv_linear_weight_scales=getattr( + layer.self_attn.qkv_proj, weight_scales_key, None + ), + # attn_output_proj + attn_output_proj_linear_smooth_scales=getattr( + layer.self_attn.o_proj, smooth_scales_key, None + ), + attn_output_proj_linear_weight_scales=getattr( + layer.self_attn.o_proj, weight_scales_key, None + ), + # mlp_linear1 + mlp_linear1_smooth_scales=getattr( + layer.mlp.gate_up_proj, smooth_scales_key, None + ), + mlp_linear1_weight_scales=getattr( + layer.mlp.gate_up_proj, weight_scales_key, None + ), + # mlp_linear2 + mlp_linear2_smooth_scales=getattr( + layer.mlp.down_proj, smooth_scales_key, None + ), + mlp_linear2_weight_scales=getattr( + layer.mlp.down_proj, weight_scales_key, None + ), + # other + activation_dtype=activation_dtype, + quant_mode=quant_mode + ) + + if params.mlp_linear2_smooth_scales is not None: + params.mlp_activation = ixff.silu_and_mul_smoothquant + + return params + + raise RuntimeError(f"unsupported quantized mode, got {quant_mode}.") + + +def create_vllm_llama_decoder_layer( + layer: torch.nn.Module, + model_id, + layer_idx, + enable_overlap=True, + group=None, + quant_config=None, + activation_dtype: Optional[torch.dtype] = None, +) -> torch.nn.Module: + """ + :param layer: vLLM 中 LLaMa 的 DecoderLayer + :param model_id: 模型的 id，可以使用 id(model) 获取 + :param layer_idx: layer 的索引 + :param enable_overlap: 是否其中 overlap + :param group: Communication group + :param quant_config: vLLM 中量化的配置 + :param activation_dtype: 在使用 w8a8 的 gemm 时，需要通过该参数去决定输出的类型 + :return: vLLM 中 LLaMa 的 DecoderLayer + """ + # 1. 如果是 overlap 不支持的量化类型 或者 不启用 overlap，那么直接返回 layer + quant_mode = None if quant_config is None else quant_config.get_name() + is_supported_quant_mode = is_vllm_supported_quant_mode(quant_mode=quant_mode) + + if ( + not enable_overlap + or not config.IXFORMER_ENABLE_OVERLAP_COMM + or not is_supported_quant_mode + ): + return layer + + # 2. 如果不是不是多卡推理，那么直接返回 layer + if hasattr(group, "device_group"): + group = group.device_group + + if not dist.is_initialized() or ixfd.get_world_size(group) < 2: + return layer + + # 3. 是否使用 rotary enmedding + rotary_embedding = None + if ( + hasattr(layer.self_attn, "postion_embedding") + and layer.self_attn.postion_embedding != "ALIBI" + ) or (not hasattr(layer.self_attn, "postion_embedding")): + rotary_embedding = layer.self_attn.rotary_emb + + # 4. 创建 Overlap 的 DecoderLayer + protocol_cls = get_vllm_llama_decoder_layer_protocol_cls( + layer=layer, quant_mode=quant_mode + ) + + params = create_vllm_llama_decoder_layer_params( + layer=layer, quant_mode=quant_mode, activation_dtype=activation_dtype + ) + + num_chunks = 2 + overlap_layer = protocol_cls( + model_id=model_id, + layer_idx=layer_idx, + attn_q_size=layer.self_attn.q_size, + attn_kv_size=layer.self_attn.kv_size, + params=params, + max_num_chunks=num_chunks, + ) + + # 5. 替换 layer 的 forward + def forward( + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_cache: KVCache, + input_metadata, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, OptionalTensor]: + """ + :param positions: 会被用在 RotaryEmbedding 中 + :param hidden_states: Shape[Batch * SeqLen, HiddenSize] + :param kv_cache: 会被使用在 Attention 中 + :param input_metadata: 会被使用在 Attention 中 + :param residual: 残差 + :return: HiddenStates, Residual + """ + return overlap_layer.forward( + self_attn=layer.self_attn.attn, + positions=positions, + hidden_states=hidden_states, + residual=residual, + num_chunks=num_chunks, + group=group, + self_attn_kwargs={ + "kv_cache": kv_cache, + "rotary_embedding": rotary_embedding, + "attn_metadata": input_metadata, + }, + ) + + layer.forward = forward + return layer diff --git a/ixformer_sdk/inference/overlap/moe_reduce_allreduce_ln_linear_overlap.py b/ixformer_sdk/inference/overlap/moe_reduce_allreduce_ln_linear_overlap.py new file mode 100644 index 00000000..b59f2a33 --- /dev/null +++ b/ixformer_sdk/inference/overlap/moe_reduce_allreduce_ln_linear_overlap.py @@ -0,0 +1,190 @@ +import dataclasses +import math +from typing import Optional, Tuple + +import ixformer.distributed as ixfd +import ixformer.functions as F +import torch +import torch.distributed as dist +from ixformer.distributed.overlap_comm import SplitOverlapComm + +from ixformer.core import config as ixff_config + + +@dataclasses.dataclass +class MoeReduceAllReduceLnQkvLinearParams: + # ============================== + # MOE Reduce Sum + # ============================== + + # shape: [Batch * SeqLen, TopK, HiddenSize], dtype: float16 or bfloat16 + input: torch.Tensor + + # shape: [Batch * SeqLen, TopK], dtype: float32 + topk_weight: Optional[torch.Tensor] + + # ============================== + # Ln + # ============================== + + # shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16 + residual: torch.Tensor + + # shape: [HiddenSize], dtype: float16 or bfloat16 + ln_weight: torch.Tensor + + # shape: [HiddenSize], dtype: float16 or bfloat16 + ln_bias: torch.Tensor + ln_eps: float + + # ============================== + # QkvLinear + # ============================== + + # shape: [(NumHeads + 2 * NumKvHeads) * HeadDim / TP, HiddenSize], dtype: float16 or bfloat16 + qkv_weight: torch.Tensor + + # shape: [(NumHeads + 2 * NumKvHeads) * HeadDim / TP] + qkv_weight_scale: torch.Tensor + + # shape: [Batch * SeqLen, (NumHeads + 2 * NumKvHeads) * HeadDim / TP], dtype: float16 or bfloat16 + qkv_out: torch.Tensor + + +class MoeReduceSumAllReduceLnQkvLinearOverlap(SplitOverlapComm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.allreduce_end_events = [torch.cuda.Event() for _ in range(self.num_chunks)] + + def start_qkv_linear(self, chunk_idx): + compute_stream = self._compute_streams[chunk_idx % self.num_compute_streams] + compute_stream.wait_event(self.allreduce_end_events[chunk_idx]) + + def compute(self, params: MoeReduceAllReduceLnQkvLinearParams, split_ratio=0.5): + input_chunk_sizes = [int(math.ceil(params.input.shape[0] * split_ratio))] + input_chunk_sizes.append(params.input.shape[0] - input_chunk_sizes[0]) + + input_chunks = list( + torch.split_with_sizes(params.input, input_chunk_sizes, dim=0) + ) + topk_weight_chunks = list( + torch.split_with_sizes(params.topk_weight, input_chunk_sizes, dim=0) + ) + out_chunks = [] + + for chunk_idx in range(len(input_chunks)): + with self.compute_stream_context(chunk_idx): + out_chunks.append( + F.moe_output_reduce_sum( + input_chunks[chunk_idx], + topk_weight=topk_weight_chunks[chunk_idx], + ) + ) + + self.start_comm(chunk_idx) + + ixfd.all_reduce( + out_chunks[chunk_idx], + async_op=True, + group=self.comm_group, + use_comm_stream=True, + ) + + self.allreduce_end_events[chunk_idx].record(self._comm_stream) + + residual_chunk_sizes = [int(params.residual.shape[0] * split_ratio)] + residual_chunk_sizes.append(params.residual.shape[0] - residual_chunk_sizes[0]) + residual_chunks = torch.split_with_sizes( + params.residual, residual_chunk_sizes, dim=0 + ) + + qkv_out_chunk_sizes = [int(params.qkv_out.shape[0] * split_ratio)] + qkv_out_chunk_sizes.append(params.qkv_out.shape[0] - qkv_out_chunk_sizes[0]) + qkv_out_chunks = torch.split_with_sizes( + params.qkv_out, qkv_out_chunk_sizes, dim=0 + ) + + for chunk_idx in range(len(input_chunks)): + self.start_qkv_linear(chunk_idx) + with self.compute_stream_context(chunk_idx): + ( + i8_hidden_states, + residual, + i_scales, + ) = F.residual_layer_norm_dynamic_int8( + input=out_chunks[chunk_idx], + residual=residual_chunks[chunk_idx], + weight=params.ln_weight, + bias=params.ln_bias, + eps=params.ln_eps, + ) + + F.w8a8( + i8_hidden_states, + params.qkv_weight, + i_scales, + params.qkv_weight_scale, + output=qkv_out_chunks[chunk_idx], + ) + + return params.qkv_out, params.residual + + +_moe_reduce_with_allreduce_overlap = None + + +def moe_reduce_sum_allreduce_ln_qkv_linear( + params: MoeReduceAllReduceLnQkvLinearParams, + enable_overlap=False, + comm_group=None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + MOE Reduce Sum + AllReduce + LayerNorm + QkvLinear + + Args: + params: fused operator params + enable_overlap: whether enable overlap + comm_group: communication group + Returns: + QkvLinearOutput: [Batch * SeqLen, (NumHeads + 2 * NumKvHeads) * HeadDim / TP], dtype: float16 or bfloat16 + Residual: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16 + """ + + global _moe_reduce_with_allreduce_overlap + if _moe_reduce_with_allreduce_overlap is None: + _moe_reduce_with_allreduce_overlap = ( + MoeReduceSumAllReduceLnQkvLinearOverlap.dispatcher( + num_chunks=2, comm_group=comm_group + ).forward + ) + + if ( + enable_overlap + and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM + and dist.is_initialized() + and dist.get_world_size(comm_group) > 1 + and params.input.shape[0] > 1 + ): + return _moe_reduce_with_allreduce_overlap(params, split_ratio=0.5) + + out = F.moe_output_reduce_sum(params.input, topk_weight=params.topk_weight) + ixfd.all_reduce(out, async_op=True, group=comm_group) + + i8_hidden_states, residual, i_scales = F.residual_layer_norm_dynamic_int8( + input=out, + residual=params.residual, + weight=params.ln_weight, + bias=params.ln_bias, + eps=params.ln_eps, + ) + + out = F.w8a8( + i8_hidden_states, + params.qkv_weight, + i_scales, + params.qkv_weight_scale, + output=params.qkv_out, + ) + + return out, residual diff --git a/ixformer_sdk/inference/overlap/overlap_comm.py b/ixformer_sdk/inference/overlap/overlap_comm.py new file mode 100644 index 00000000..52808d19 --- /dev/null +++ b/ixformer_sdk/inference/overlap/overlap_comm.py @@ -0,0 +1,250 @@ +import enum +from typing import Optional + +import ixformer.functions as ixff +import torch +from ixformer.inference.overlap.linear_mlp_overlap_comm import ( + LinearMLPOverlapComm, + LinearMLPOverlapCommHook, +) + + +def get_overlap_linear_method(layer): + if hasattr(layer, "_overlap_comm_gemm_fn"): + return layer._overlap_comm_gemm_fn + + if layer.linear_weights["weight"].itemsize == 2: + layer._overlap_comm_gemm_fn = None + return None + + def overlap_linear_fn(input, weight, bias=None, out: torch.Tensor = None, **kwargs): + return layer.linear_method.apply_weights( + layer.linear_weights, input, output=out + ) + + layer._overlap_comm_gemm_fn = overlap_linear_fn + return overlap_linear_fn + + +class DecoderLayerOverlapComm(LinearMLPOverlapCommHook): + class HookStage(enum.IntEnum): + kExited = 0 + kTracing = 1 + + class HookState: + def __init__(self, max_num_chunks): + self.max_num_chunks = max_num_chunks + self.stage = DecoderLayerOverlapComm.HookStage.kExited + + self.mlp_linaer2_end_events = [ + torch.cuda.Event() for _ in range(max_num_chunks) + ] + self.ln_attn_end_event = torch.cuda.Event() + + self.overlap_comm: Optional[LinearMLPOverlapComm] = None + + def is_tracing_stage(self): + return self.stage == DecoderLayerOverlapComm.HookStage.kTracing + + def enter(self, overlap_comm, chunk_idx): + self.stage = DecoderLayerOverlapComm.HookStage.kTracing + + self.overlap_comm = overlap_comm + self.mlp_linaer2_end_events[chunk_idx].record(overlap_comm._comm_stream) + + def exit(self): + self.overlap_comm = None + self.stage = DecoderLayerOverlapComm.HookStage.kExited + + def __str__(self): + return f"HookState(overlap_comm={self.overlap_comm}, stage={self.stage})" + + def __repr__(self): + return self.__str__() + + _overlap_comm_hook_state = dict() + + def __init__(self, model_id, layer_idx, max_num_chunks: int = 4): + """ + DecoderLayer 的流程： + ln_qkv: InputLayerNorm(hidden_states, [residual]) -> qkv_proj(hidden_states) -> q, k, v = split(hidden_states) -> Attention(q, k, v) + linear_mlp: AttentionOutputProj(hidden_states) -> PostLayerNorm(hidden_states) -> MLPLinear1 -> MLPActivation -> MLPLinear2 + + 其中：AttentionOutputProj 和 MLPLinear2 之后如果使用 TP，那么需要进行 AllReduce + + 通过上述流程，该类的目的是将 MLPLinear2 后的 AllReduce 和 DecoderLayer 最开始的 ln_qkv 进行 Overlap。 + 其中，第一层 DecoderLayer 不进行 ln_qkv 的 Overlap，因为在第一层之前没有通讯。 + 我们需要将第 i 层 MLPLinear2 后的通讯 和 第 i + 1 层的 ln_qkv 进行 Overlap。 + + 为了管理当前的状态和获取前一层的状态，从而设计了 DecoderLayerOverlapComm 类。 + 该类需要 model_id 来推断当前正在运行的模型，用 layer_idx 来标记每一层的开始和结束， + 以及通过 layer_idx 去获取前一层的状态。 + + 注： + - 在 call_ln_qkv_overlap 中对 Tensor 进行切分时， + 需要保持和 linear_mlp 切分的大小是一致的，否则会出现 Tensor 的数据不对应； + - 如果需要使用 ln_qkv 进行 Overlap，那么必须使用该类的 linear_mlp 去替换 linear_mlp_overlap + + :param model_id: 模型的 id，可以使用 id(model) 去设置 + :param layer_idx: layer 的索引，注意，需要从 0 到 NumLayers 的顺序去完成构造 + :param max_num_chunks: 最大能进行切分的次数 + """ + self._model_id = model_id + self._layer_idx = layer_idx + self._max_num_chunks = max_num_chunks + + self._state = self.HookState(max_num_chunks) + self._overlap_comm_hook_state[(model_id, layer_idx)] = self._state + self._prev_layer_state = ( + None + if layer_idx == 0 + else self._overlap_comm_hook_state[(model_id, layer_idx - 1)] + ) + + @property + def model_id(self): + return self._model_id + + @property + def layer_idx(self): + return self._layer_idx + + @property + def max_num_chunks(self): + return self._max_num_chunks + + @property + def state(self) -> "DecoderLayerOverlapComm.HookState": + return self._state + + @property + def prev_layer_state(self) -> "DecoderLayerOverlapComm.HookState": + return self._prev_layer_state + + def is_ln_qkv_overlap(self): + return not ( + self.layer_idx == 0 + or not self.prev_layer_state.is_tracing_stage() + or self.prev_layer_state.overlap_comm is None + ) + + def ln_qkv(self, hidden_states, residual, ln_layer, qkv_layer, out_last_dim): + """ + :param hidden_state: shape[Batch * SeqLen, HiddenSize] + :param residual: shape[Batch * SeqLen, HiddenSize] + :param ln_layer: torch.nn.Module or Function(hidden_state, residual=None) + :param qkv_layer: vllm.QKVParallelLinear + :param out_last_dim: qkv_layer 输出 Tensor 的最后一个维度 + :return: qkv, residual + """ + if self.is_ln_qkv_overlap(): + qkv, residual = self.call_ln_qkv_overlap( + hidden_states, residual, ln_layer, qkv_layer, out_last_dim + ) + else: + qkv, residual = self.call_ln_qkv( + hidden_states, residual, ln_layer, qkv_layer + ) + + return qkv, residual + + def call_ln_qkv_overlap( + self, hidden_states, residual, ln_layer, qkv_layer, out_last_dim + ): + if hidden_states.ndim != 2: + raise RuntimeError( + f"Expected 2-dim for hidden state, but got {hidden_states.ndim}." + ) + + num_chunks = self.prev_layer_state.overlap_comm.num_chunks + overlap_comm: LinearMLPOverlapComm = self.prev_layer_state.overlap_comm + + if num_chunks > self.max_num_chunks: + raise RuntimeError( + f"The layer is not support more than {self.max_num_chunks}, got {num_chunks}." + ) + + hidden_state_chunks = list(torch.chunk(hidden_states, num_chunks, dim=0)) + if residual is None: + residual = hidden_states + residual_chunks = [None] * num_chunks + else: + residual_chunks = torch.chunk(residual, num_chunks, dim=0) + + out = torch.empty( + (hidden_states.shape[0], out_last_dim), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + out_chunks = list(torch.chunk(out, num_chunks, dim=0)) + + for chunk_idx, (hidden_state_chunk, residual_chunk, out_chunk) in enumerate( + zip(hidden_state_chunks, residual_chunks, out_chunks) + ): + overlap_comm._compute_streams[ + chunk_idx % overlap_comm.num_compute_streams + ].wait_event(self.prev_layer_state.mlp_linaer2_end_events[chunk_idx]) + with overlap_comm.compute_stream_context(chunk_idx): + self.call_ln_qkv( + hidden_state_chunk, + residual_chunk, + ln_layer, + qkv_layer, + chunk_idx, + use_limited_gemm=chunk_idx != (num_chunks - 1), + out=out_chunk, + overlap_comm=overlap_comm, + ) + + self.prev_layer_state.exit() + overlap_comm.stop_overlap() + return out, residual + + def call_ln_qkv( + self, + hidden_state, + residual, + ln_layer, + qkv_layer, + chunk_idx=0, + use_limited_gemm=False, + out=None, + overlap_comm: LinearMLPOverlapComm = None, + ): + if residual is None: + residual = hidden_state + if ln_layer is not None: + hidden_state = ln_layer(hidden_state) + else: + hidden_state, residual = ln_layer(hidden_state, residual) + + if out is None: + qkv, _ = qkv_layer(hidden_state) + else: + gemm_method = get_overlap_linear_method(qkv_layer) + qkv = overlap_comm.gemm_dispatcher( + chunk_idx=chunk_idx, + chunk_input=hidden_state, + weight=qkv_layer.linear_weights["weight"], + chunk_out=out, + user_gemm_method=gemm_method, + use_limited_gemm=use_limited_gemm, + ) + + return qkv, residual + + def linear_mlp(self, *args, **kwargs): + """ref: linear_mlp_overlap""" + return ixff.linear_mlp_overlap( + *args, **kwargs, mlp_linear2_finished_callback=self.on_mlp_linear2_finished + ) + + def on_mlp_linear2_finished( + self, + overlap_comm: LinearMLPOverlapComm, + num_chunks, + chunk_idx, + hidden_states_chunk, + residual_chunk, + ): + self.state.enter(overlap_comm, chunk_idx) diff --git a/ixformer_sdk/inference/overlap/w8a8_allreduce.py b/ixformer_sdk/inference/overlap/w8a8_allreduce.py new file mode 100644 index 00000000..157fffb3 --- /dev/null +++ b/ixformer_sdk/inference/overlap/w8a8_allreduce.py @@ -0,0 +1,154 @@ +import math +from typing import Optional + +import ixformer.distributed as ixfd +import ixformer.functions as F +import torch +import torch.distributed as dist +from ixformer.distributed.overlap_comm import SplitOverlapComm + +from ixformer.core import config as ixff_config + +__all__ = ["w8a8_allreduce"] + + +class W8A8AllReduceOverlap(SplitOverlapComm): + def compute( + self, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + bias: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + format: str = "TN", + out_dtype: torch.dtype = None, + comm_group=None, + split_ratio=0.5, + ): + # compute the chunk size of input + input_chunk_sizes = [int(math.ceil(input.shape[0] * split_ratio))] + input_chunk_sizes.append(input.shape[0] - input_chunk_sizes[0]) + + # split input and input_scale + input_chunks = list(torch.split_with_sizes(input, input_chunk_sizes, dim=0)) + input_scale_chunks = torch.split( + input_scale, + input_chunk_sizes, + ) + + # create output and split it + if output is None: + if out_dtype is None: + raise RuntimeError( + "w8a8 gemm need out_dtype argument when output is none." + ) + output = torch.empty( + (input.shape[:-1] + (weight.shape[0],)), + dtype=out_dtype, + device=input.device, + ) + out_chunks = torch.split(output, input_chunk_sizes) + + # overlap gemm and allreduce + for chunk_idx in range(len(input_chunks)): + # submit gemm kernel into compute stream + with self.compute_stream_context(chunk_idx): + F.w8a8( + input=input_chunks[chunk_idx], + weight=weight, + i_scales=input_scale_chunks[chunk_idx], + w_scales=weight_scale, + bias=bias, + output=out_chunks[chunk_idx], + format=format, + persistent=chunk_idx != 0, + ) + + # recode compute stream and wait gemm + self.start_comm(chunk_idx) + + # submit allreduce kernel into communication stream by set use_comm_stream to true + ixfd.all_reduce( + out_chunks[chunk_idx], + async_op=True, + group=self.comm_group, + use_comm_stream=True, + ) + + return output + + +_w8a8_allreduce_overlap = None + + +def w8a8_allreduce( + enable_overlap: bool, + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + bias: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + format: str = "TN", + out_dtype: torch.dtype = None, + comm_group=None, + split_ratio=0.5, +) -> torch.Tensor: + """ + Gemm(w8a8) + AllReduce + + Args: + enable_overlap: whether enable gemm and allreduce overlap + input: shape: [M, K], dtype: int8, linear input + weight: shape: [N, K], dtype: int8, linear weight + input_scale: shape: [M], dtype: float32, quantized scale of input + weight_scale: shape: [N], dtype: float32, quantized scale of weight + bias: shape: [N], dtype: float16 or bfloat16, linear bias + output: shape: [M, N], dtype: float16 or bfloat16, allreduce output + format: options include TN, NN and NT + out_dtype: use the argument to decide to the dtype of output when output is None + comm_group: communication group + split_ratio: split the ratio of input.shape[0] when using overlap, range: (0, 1), + it will affect area of the overlap for gemm and allreduce. + Returns: output + """ + if ( + enable_overlap + and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM + and dist.is_initialized() + and dist.get_world_size(comm_group) > 1 + and input.shape[0] > 1 + ): + global _w8a8_allreduce_overlap + if _w8a8_allreduce_overlap is None: + _w8a8_allreduce_overlap = W8A8AllReduceOverlap.dispatcher( + num_chunks=2, comm_group=comm_group + ).forward + return _w8a8_allreduce_overlap( + input=input, + weight=weight, + input_scale=input_scale, + weight_scale=weight_scale, + bias=bias, + output=output, + format=format, + out_dtype=out_dtype, + split_ratio=split_ratio, + ) + + out = F.w8a8( + input=input, + weight=weight, + i_scales=input_scale, + w_scales=weight_scale, + bias=bias, + output=output, + format=format, + out_dtype=out_dtype, + ) + + if dist.get_world_size() > 1: + ixfd.all_reduce(out, op=ixfd.ReduceOp.SUM, async_op=True, group=comm_group) + + return out diff --git a/ixformer_sdk/testing/__init__.py b/ixformer_sdk/testing/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/testing/memory_monitor.py b/ixformer_sdk/testing/memory_monitor.py new file mode 100644 index 00000000..07d97ed3 --- /dev/null +++ b/ixformer_sdk/testing/memory_monitor.py @@ -0,0 +1,37 @@ +from psutil import Process + + +def get_current_memory(pid=None): + return Process(pid).memory_full_info() + + +class MemoryMonitorContext(object): + def __init__(self, pid=None): + self._pid = pid + + self.reset() + + def __enter__(self): + self._enter_memory = self._get_used_memory() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self._exit_memory = self._get_used_memory() + return self + + def _get_used_memory(self) -> float: + mem = get_current_memory(self._pid) + return mem.uss + + def reset(self): + self._enter_memory = -1 + self._exit_memory = -1 + + def delta(self) -> float: + if self._enter_memory < 0: + raise RuntimeError("Please using context manager to wrap your code.") + + if self._exit_memory < 0: + return self._get_used_memory() - self._enter_memory + + return self._exit_memory - self._enter_memory diff --git a/ixformer_sdk/train/__init__.py b/ixformer_sdk/train/__init__.py new file mode 100644 index 00000000..e7d4dc00 --- /dev/null +++ b/ixformer_sdk/train/__init__.py @@ -0,0 +1 @@ +from .functions import * \ No newline at end of file diff --git a/ixformer_sdk/train/functions/__init__.py b/ixformer_sdk/train/functions/__init__.py new file mode 100644 index 00000000..97e3b8f9 --- /dev/null +++ b/ixformer_sdk/train/functions/__init__.py @@ -0,0 +1,12 @@ +from .cross_entropy_loss import * +from .fused_rope import * +from .geglu import * +from .gelu import * +from .layernorm import * +from .linear import * +from .matmul import * +from .residual_bias import * +from .residual_bias_ln import * +from .rms_norm import * +from .swiglu import * +from .group_norm import * diff --git a/ixformer_sdk/train/functions/cross_entropy_loss.py b/ixformer_sdk/train/functions/cross_entropy_loss.py new file mode 100644 index 00000000..ab87b72b --- /dev/null +++ b/ixformer_sdk/train/functions/cross_entropy_loss.py @@ -0,0 +1,239 @@ +from typing import Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function +from torch.nn import init +from torch.nn.parameter import Parameter + +__all__ = ["vocab_parallel_cross_entropy"] + + +class _VocabParallelCrossEntropyCustom(Function): + @staticmethod + def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0): + device = vocab_parallel_logits.device + xnumel = vocab_parallel_logits.shape[0] + rnumel = vocab_parallel_logits.shape[-1] + exp_logits = torch.empty( + (xnumel, 1, rnumel), device=device, dtype=torch.float32 + ) + masked_target_1d = torch.empty((xnumel,), device=device, dtype=torch.int32) + loss = torch.empty((xnumel, 1), device=device, dtype=torch.float32) + + ops.train.cross_entropy_loss_forward( + vocab_parallel_logits, target.int(), exp_logits, masked_target_1d, loss + ) + + vocab_size = exp_logits.size(-1) + if label_smoothing > 0: + """ + We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth. + = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt}) + = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i + = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K + From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py + """ + assert 1.0 > label_smoothing > 0.0 + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + + # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs. + log_probs = torch.log(exp_logits) + mean_log_probs = log_probs.mean(dim=-1) + loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs + + ctx.label_smoothing, ctx.vocab_size = label_smoothing, vocab_size + + # Store softmax, target-mask and masked-target for backward pass. + ctx.save_for_backward(exp_logits, masked_target_1d) + + return loss + + @staticmethod + def backward(ctx, grad_output): + + # Retreive tensors from the forward path. + softmax, masked_target_1d = ctx.saved_tensors + label_smoothing, vocab_size = ctx.label_smoothing, ctx.vocab_size + + # All the inputs have softmax as thier gradient. + grad_input = softmax + # For simplicity, work with the 2D gradient. + partition_vocab_size = softmax.size()[-1] + grad_2d = grad_input.view(-1, partition_vocab_size) + + # Add the gradient from matching classes. + arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device) + + softmax_update = 1.0 + + if label_smoothing > 0: + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + grad_2d[arange_1d, masked_target_1d] -= (1.0 - smoothing) * softmax_update + average_grad = 1 / vocab_size + grad_2d[arange_1d, :] -= smoothing * average_grad + else: + grad_2d[arange_1d, masked_target_1d] -= softmax_update + + # Finally elementwise multiplication with the output gradients. + grad_input = torch.mul(grad_input, grad_output.unsqueeze(dim=-1)) + + return grad_input, None, None + + +class _VocabParallelCrossEntropy(Function): + @staticmethod + def forward( + ctx, + vocab_parallel_logits, + target, + label_smoothing=0.0, + vocab_start_index=0, + vocab_end_index=320000, + group=None, + ): + + # Maximum value along vocab dimension across all GPUs. + logits_max = torch.max(vocab_parallel_logits, dim=-1)[0] + torch.distributed.all_reduce( + logits_max, op=torch.distributed.ReduceOp.MAX, group=group + ) + # Subtract the maximum value. + vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1) + + # Get the partition's vocab indecies + partition_vocab_size = vocab_parallel_logits.size()[-1] + + # Create a mask of valid vocab ids (1 means it needs to be masked). + target_mask = (target < vocab_start_index) | (target >= vocab_end_index) + masked_target = target.clone() - vocab_start_index + masked_target[target_mask] = 0 + + # Get predicted-logits = logits[target]. + # For Simplicity, we convert logits to a 2-D tensor with size + # [*, partition-vocab-size] and target to a 1-D tensor of size [*]. + logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size) + masked_target_1d = masked_target.view(-1) + arange_1d = torch.arange( + start=0, end=logits_2d.size()[0], device=logits_2d.device + ) + predicted_logits_1d = logits_2d[arange_1d, masked_target_1d] + predicted_logits_1d = predicted_logits_1d.clone().contiguous() + predicted_logits = predicted_logits_1d.view_as(target) + predicted_logits[target_mask] = 0.0 + # All reduce is needed to get the chunks from other GPUs. + torch.distributed.all_reduce( + predicted_logits, + op=torch.distributed.ReduceOp.SUM, + group=group, + ) + + # Sum of exponential of logits along vocab dimension across all GPUs. + exp_logits = vocab_parallel_logits + torch.exp(vocab_parallel_logits, out=exp_logits) + sum_exp_logits = exp_logits.sum(dim=-1) + torch.distributed.all_reduce( + sum_exp_logits, + op=torch.distributed.ReduceOp.SUM, + group=group, + ) + + # Loss = log(sum(exp(logits))) - predicted-logit. + loss = torch.log(sum_exp_logits) - predicted_logits + + # Normalize and optionally smooth logits + exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1)) + + vocab_size = exp_logits.size(-1) + if label_smoothing > 0: + """ + We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth. + = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt}) + = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i + = (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i + = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K + From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py + """ + assert 1.0 > label_smoothing > 0.0 + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + + # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs. + log_probs = torch.log(exp_logits) + mean_log_probs = log_probs.mean(dim=-1) + loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs + + ctx.label_smoothing, ctx.vocab_size = label_smoothing, vocab_size + + # Store softmax, target-mask and masked-target for backward pass. + ctx.save_for_backward(exp_logits, target_mask, masked_target_1d) + + return loss + + @staticmethod + def backward(ctx, grad_output): + + # Retreive tensors from the forward path. + softmax, target_mask, masked_target_1d = ctx.saved_tensors + label_smoothing, vocab_size = ctx.label_smoothing, ctx.vocab_size + + # All the inputs have softmax as thier gradient. + grad_input = softmax + # For simplicity, work with the 2D gradient. + partition_vocab_size = softmax.size()[-1] + grad_2d = grad_input.view(-1, partition_vocab_size) + + # Add the gradient from matching classes. + arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device) + + softmax_update = 1.0 - target_mask.view(-1).float() + + if label_smoothing > 0: + smoothing = label_smoothing * vocab_size / (vocab_size - 1) + grad_2d[arange_1d, masked_target_1d] -= (1.0 - smoothing) * softmax_update + average_grad = 1 / vocab_size + grad_2d[arange_1d, :] -= smoothing * average_grad + else: + grad_2d[arange_1d, masked_target_1d] -= softmax_update + + # Finally elementwise multiplication with the output gradients. + grad_input.mul_(grad_output.unsqueeze(dim=-1)) + + return grad_input, None, None, None, None, None + + +def vocab_parallel_cross_entropy( + vocab_parallel_logits: torch.Tensor, + target: torch.Tensor, + label_smoothing: float = 0.0, + world_size: int = 1, + vocab_start_index: int = 0, + vocab_end_index: int = 320000, + group=None, +): + """ + 参数说明: + 目前只支持batch_size = 1 的情况,当batch_size >1时,计算不正确 + Args: + vocab_parallel_logits: shape : [seq_len,1,vocal_size] dtype : torch.bfloat16,torch.float,torch.half + target: shape : [seq_len,1] dtype : torch.int64 + group: TP 并行组 + return: + loss: shape : [seq_len,1] dtype : torch.float32 + + """ + if world_size == 1: + return _VocabParallelCrossEntropyCustom.apply( + vocab_parallel_logits, target, label_smoothing + ) + else: + return _VocabParallelCrossEntropy.apply( + vocab_parallel_logits, + target, + label_smoothing, + vocab_start_index, + vocab_end_index, + group, + ) diff --git a/ixformer_sdk/train/functions/fused_rope.py b/ixformer_sdk/train/functions/fused_rope.py new file mode 100644 index 00000000..198140d8 --- /dev/null +++ b/ixformer_sdk/train/functions/fused_rope.py @@ -0,0 +1,214 @@ +from typing import List, Tuple, Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function + +# adding by xuelu.peng 20240417 +# from https://github.com/NVIDIA/apex/blob/master/apex/transformer/functional/fused_rope.py#L59 +__all__ = ["fused_apply_rotary_pos_emb", "fused_apply_split_rotary_pos_emb", "fused_apply_rotary_pos_emb_cache"] + + +class FusedRoPEFunc(Function): + """ + Fused RoPE function + + This implementation assumes the input tensor to be in `sbhd` format and the RoPE tensor to be + of shape (s, 1, 1, d). It accepts arbitrary memory layouts to avoid the expensive + `.contiguous()` calls, thus it may not achieve the best memory access pattern. + """ + + @staticmethod + def forward( + ctx, + t: torch.Tensor, + freqs: torch.Tensor, + transpose_output_memory: bool = False, + ) -> torch.Tensor: + # assert transpose_output_memory == False + output = ops.train.fused_rope_forward(t, freqs, transpose_output_memory) + ctx.save_for_backward(freqs) + ctx.transpose_output_memory = transpose_output_memory + + return output + + @staticmethod + def backward( + ctx, grad_output: torch.Tensor + ) -> Tuple[Union[torch.Tensor, None], ...]: + + (freqs,) = ctx.saved_tensors + grad_input = ops.train.fused_rope_backward( + grad_output, freqs, ctx.transpose_output_memory + ) + return grad_input, None, None + + +class FusedFluxRoPEFunc(Function): + """ + Fused FluxRoPE function + + This implementation assumes the input tensor to be in `bshd` format and the RoPE tensor to be + of shape (s, d), and output shape is the same as input shape. + """ + + @staticmethod + def forward( + ctx, + t: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + imp_mode : int = 1 + ) -> torch.Tensor: + # assert transpose_output_memory == False + output = ops.train.fused_rope_forward_cached(t, cos, sin, imp_mode) + ctx.save_for_backward(cos, sin) + return output + + @staticmethod + def backward( + ctx, grad_output: torch.Tensor + ) -> Tuple[Union[torch.Tensor, None], ...]: + (cos, sin) = ctx.saved_tensors + grad_input = ops.train.fused_rope_backward_cached( + grad_output, cos, sin + ) + return grad_input, None, None, None + + + +def fused_apply_rotary_pos_emb( + t: torch.Tensor, + freqs: torch.Tensor, + transpose_output_memory: bool = False, +) -> torch.Tensor: + """Apply rotary positional embedding to input tensor T in `sbhd` format, where + s: sequence length + b: batch size + h: head num + d: dim of each head + + Args: + t (Tensor): Input tensor T is of shape [s, b, h, d], dtype : torch.float32, torch.half + freqs (Tensor): Rotary Positional embedding tensor freq is of shape [s, 1, 1, d] and + `float` dtype + transpose_output_memory (bool): Default to False. Whether to transpose the 's' and 'b' + dimension of the output's underlying memory format. This is very helpful when you want to + get a contiguous tensor after calling `output.transpose(0, 1)`. + + Returns: + Tensor: The input tensor after applying RoPE + """ + return FusedRoPEFunc.apply(t, freqs, transpose_output_memory) + + +def fused_apply_rotary_pos_emb_cache( + t: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + imp_mode: int = 1, +) -> torch.Tensor: + """Apply rotary positional embedding to input tensor T in `bshd` format, where + s: sequence length + b: batch size + h: head num + d: dim of each head + + Args: + t (Tensor): Input tensor T is of shape [b, s, h, d], dtype : torch.float32, torch.half, torch.bfloat16 + cos/sin (Tensor): Rotary Positional embedding tensor freq is of shape [s, d] and + `float` dtype + imp_mode (bool): Default to 1. 1 for flux/cogvideox/hunyuan-dit, img_mode=0 for Stable Audio. For now, only img_mode = 1 is supported. + + Returns: + Tensor: The input tensor after applying RoPE + """ + return FusedFluxRoPEFunc.apply(t, cos, sin, imp_mode) + + +class FusedSplitRoPEFunc(torch.autograd.Function): + """ + Fused Split and RoPE function + + This implementation assumes the input tensor to be in `sbh3d` format and the RoPE tensor to be + of shape (s, 1, 1, d). It accepts arbitrary memory layouts to avoid the expensive + `.contiguous()` calls, thus it may not achieve the best memory access pattern. + + input: mix_q_k_v [s,b,hn_kv,h/hn_kv+2,d] + output: output_q, output_k, output_v [s,b,h,d] + """ + @staticmethod + def forward( + ctx, + mixed_q_k_v: torch.Tensor, + freqs: torch.Tensor, + transpose_output_memory: bool = False, + ) -> torch.Tensor: + assert transpose_output_memory == False, "do not support transpose_output now" + assert mixed_q_k_v.is_contiguous() == True, "mixed_q_k_v should be contiguous in FusedSplitRoPEFunc." + + s, b, hn_kv, repplus2, d = mixed_q_k_v.size() + num_key_value_groups = repplus2-2 + + q, k, v = torch.split(mixed_q_k_v, (num_key_value_groups,1,1), dim=3) + + ctx.hn_kv = hn_kv + output_q, output_k, output_v = torch.empty_like(q).view(s,b,-1,d),torch.empty_like(q).view(s,b,-1,d),torch.empty_like(q).view(s,b,-1,d) + + ops.train.fused_split_rope_forward( + q, k, v, freqs, output_q, output_k, output_v, transpose_output_memory, hn_kv, num_key_value_groups + ) + + ctx.save_for_backward(freqs) + ctx.transpose_output_memory = transpose_output_memory + + return output_q, output_k, output_v + + @staticmethod + def backward( + ctx, grad_o_q: torch.Tensor, grad_o_k: torch.Tensor, grad_o_v: torch.Tensor + ) -> Tuple[Union[torch.Tensor, None], ...]: + # grad_o_q: [s,b,h,d] + s,b,h,d = grad_o_q.size() + + hn_kv = ctx.hn_kv + + mixed_shape = (s, b, hn_kv,(h//hn_kv+2), d) + + if hn_kv == h: + grad_mixed_q_k_v = torch.empty(mixed_shape, dtype=grad_o_q.dtype, device=grad_o_q.device,memory_format=torch.contiguous_format) # torch.empty效率比torch.zeros高 + else: + grad_mixed_q_k_v = torch.zeros(mixed_shape, dtype=grad_o_q.dtype, device=grad_o_q.device) # 支持 gqa 的情况，kernel内需要进行累加，需要把qkv的梯度置零 + grad_q, grad_k, grad_v = torch.split(grad_mixed_q_k_v.view(s,b,hn_kv,-1,d), (h//hn_kv,1,1), dim=3) + + (freqs,) = ctx.saved_tensors + ops.train.fused_split_rope_backward( + grad_o_q, grad_o_k, grad_o_v, freqs, grad_q, grad_k, grad_v, ctx.transpose_output_memory + ) + + return grad_mixed_q_k_v, None, None + +def fused_apply_split_rotary_pos_emb( + mixed_q_k_v: torch.Tensor, + freqs: torch.Tensor, + transpose_output_memory: bool = False, +) -> torch.Tensor: + """ Split mixed_q_k_v and apply rotary positional embedding to q and k in `sbhd` format, where + s: sequence length + b: batch size + h: head num + d: dim of each head + hn_kv: num head of key and value + + Args: + mixed_q_k_v (Tensor): Input tensor T is of shape [s,b,hn_kv,h/hn_kv+2,d] + freqs (Tensor): Rotary Positional embedding tensor freq is of shape [s, 1, 1, d] and + `float` dtype + transpose_output_memory (bool): Default to False. Whether to transpose the 's' and 'b' + dimension of the output's underlying memory format. This is very helpful when you want to + get a contiguous tensor after calling `output.transpose(0, 1)`. + + Returns: + Tensors: The input tensors after split and applying RoPE + """ + return FusedSplitRoPEFunc.apply(mixed_q_k_v, freqs, transpose_output_memory) diff --git a/ixformer_sdk/train/functions/geglu.py b/ixformer_sdk/train/functions/geglu.py new file mode 100644 index 00000000..defde8f1 --- /dev/null +++ b/ixformer_sdk/train/functions/geglu.py @@ -0,0 +1,43 @@ +from typing import Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = ["geglu"] + + +class GegluFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + output_shape = list(input.shape) + output_shape[-1] = output_shape[-1] // 2 + output = input.new_empty(output_shape) + ops.train.geglu_training_forward(input, output) + ctx.save_for_backward(input) + return output + + @staticmethod + def backward(ctx: FunctionCtx, grad_output): + input = ctx.saved_tensors[0] + grad_input = torch.empty_like(input) + ops.train.geglu_training_backward(input, grad_output, grad_input) + return grad_input + + +def geglu(input: "torch.Tensor"): + """ + 等价实现： + def ref_gelu_and_mul(x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + x1, x2 = x.chunk(chunks=2, dim=-1) + res = NNF.gelu(x2) * x1 + return res.to(dtype) + + Args: + input: dtype:[torch.float, torch.half, torch.bfloat16] + Returns: + output: (....,input.shape[-1] //2), dtype:[torch.float, torch.half, torch.bfloat16] + """ + return GegluFunction.apply(input) diff --git a/ixformer_sdk/train/functions/gelu.py b/ixformer_sdk/train/functions/gelu.py new file mode 100644 index 00000000..e4535ac3 --- /dev/null +++ b/ixformer_sdk/train/functions/gelu.py @@ -0,0 +1,47 @@ +from typing import List, Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = [ + "gelu", +] + + +class GeluFunction(Function): + @staticmethod + def forward( + ctx, input: torch.Tensor, in_place: bool = False, training: bool = False + ): + if training: + if in_place: + ctx.save_for_backward(input.clone()) + else: + ctx.save_for_backward(input) + if in_place: + return ops.train.gelu_forward(input, input) + else: + return ops.train.gelu_forward(input) + + @staticmethod + def backward(ctx: FunctionCtx, grad_outputs): + input = ctx.saved_tensors[0] + grad_input = ops.train.gelu_backward(input, grad_outputs) + return grad_input, None, None + + +def gelu( + input: torch.Tensor, in_place: bool = False, training: bool = False +) -> torch.Tensor: + """ + 等价实现： + torch.nn.functional.gelu + + Args: + input: dtype:[torch.float, torch.half, torch.bfloat16] + in place: bool. Whether to operate directly on the original input data. + Returns: + output: dtype:[torch.float, torch.half, torch.bfloat16] + """ + return GeluFunction.apply(input, in_place, training) diff --git a/ixformer_sdk/train/functions/group_norm.py b/ixformer_sdk/train/functions/group_norm.py new file mode 100644 index 00000000..a5042735 --- /dev/null +++ b/ixformer_sdk/train/functions/group_norm.py @@ -0,0 +1,61 @@ +import ixformer._C as ops +import torch +from torch.nn import init +from torch.nn.parameter import Parameter + + +class GN_NHWC_Func(torch.autograd.Function): + @staticmethod + def forward(ctx, X: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, G: int, eps: float, activation: str): + X_out, means, rstds = ops.train.gn_nhwc_fwd(X, weight, bias, G, eps, activation) + ctx.save_for_backward(X, weight, bias, means, rstds) + ctx.G = G + ctx.activation = activation + return X_out + + @staticmethod + def backward(ctx, dy: torch.Tensor): + dy = dy.contiguous(memory_format=torch.channels_last) + X, weight, bias, means, rstds = ctx.saved_tensors + dx, dgamma, dbeta = ops.train.gn_nhwc_bwd(dy, X, weight, bias, means, rstds, ctx.G, ctx.activation) + return dx, dgamma, dbeta, None, None, None + + +class GroupNorm_nhwc(torch.nn.GroupNorm): + def __init__(self, num_groups: int, nc: int, activation='identity', **kwargs): + super().__init__(num_groups, nc, **kwargs) + assert activation in {'identity', 'silu', 'relu', 'gelu', 'gelu_tanh'} + if activation == 'identity': + self.activation = 0 + if activation == 'relu': + self.activation = 1 + if activation == 'silu': + self.activation = 2 + if activation == 'gelu': + self.activation = 3 + if activation == 'gelu_tanh': + self.activation = 4 + + @torch._dynamo.disable + def forward(self, x): + #print(x.shape, self.num_channels) + if len(x.size()) == 3: + N, C, L = x.shape + elif len(x.size()) == 4: + N, C, H, W = x.shape + else: + raise ValueError + G = self.num_groups + + #if C // G > 512: + # raise ValueError(f'Error in fwd for X.shape={x.shape}, G={G}: C // G = {C // G} which is greater than 512. This input is not supported.') + + #if H * W % 8 != 0: + # raise ValueError(f'Error in fwd for X.shape={x.shape}, G={G}: H * W is not a multiple of 8. This input is not supported.') + + if self.affine: + return GN_NHWC_Func.apply(x, self.weight, self.bias, self.num_groups, self.eps, self.activation) + else: + w = torch.ones((self.num_channels,), device=x.device, dtype=x.dtype) + b = torch.zeros((self.num_channels,), device=x.device, dtype=x.dtype) + return GN_NHWC_Func.apply(x, w, b, self.num_groups, self.eps, self.activation) diff --git a/ixformer_sdk/train/functions/layernorm.py b/ixformer_sdk/train/functions/layernorm.py new file mode 100644 index 00000000..4a0de167 --- /dev/null +++ b/ixformer_sdk/train/functions/layernorm.py @@ -0,0 +1,98 @@ +from typing import List, Tuple, Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = ["layernorm"] + + +class LayerNormFunction(Function): + @staticmethod + def forward( + ctx, + input: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + output: torch.Tensor, + normalized_shape=None, + training: bool = False, + ): + + if ln_weight is None or ln_bias is None: + raise NotImplementedError() + # normalized_shape 需要是list或者tuple，并且不能为空 + if normalized_shape == None: + norm_size = ln_weight.size(-1) + else: + norm_size = 1 + if isinstance(normalized_shape, int): + norm_size = normalized_shape + normalized_shape = [normalized_shape] + + elif ( + isinstance(normalized_shape, list) + or isinstance(normalized_shape, tuple) + ) and len(normalized_shape) >= 1: + for i in normalized_shape: + norm_size = i * norm_size + else: + raise f"layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints, not {type(normalized_shape)}" + if norm_size != ln_weight.size(-1): + raise f"layer_norm(): argument 'norm_size' must == ln_weight.size(-1)" + if output is None: + output = torch.empty_like(input) + if training: + mean_size = input.numel() // norm_size + + input_hat = torch.empty_like(input) + rstd = torch.empty([mean_size], dtype=input.dtype, device=input.device) + ops.train.layernorm_training_forward( + input, ln_weight, ln_bias, output, input_hat, rstd + ) + ctx.norm_size = norm_size + ctx.save_for_backward(input_hat, rstd, ln_weight) + else: + ops.train.layernorm_forward(input, ln_weight, ln_bias, output) + return output + + @staticmethod + # def backward(ctx: FunctionCtx, grad_output, dh, dr): + def backward(ctx: FunctionCtx, grad_output): + input_hat, rstd, ln_weight = ctx.saved_tensors + + grad_input = torch.empty_like(input_hat) + grad_weight = torch.empty_like(ln_weight) + grad_bias = torch.empty_like(ln_weight) + ops.train.layernorm_weightbias_backward( + input_hat, grad_output, grad_weight, grad_bias + ) + ops.train.layernorm_input_backward( + input_hat, rstd, grad_output, ln_weight, grad_input + ) + return grad_input, grad_weight, grad_bias, None, None, None + + +def layernorm( + input: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + normalized_shape=None, + output: torch.Tensor = None, + training: bool = False, +): + """ + 等价实现： + torch.nn.functional.layer_norm( input, normalized_shape, ln_weight, ln_bias, eps=0.000001) + Arguments: + input: (batch_count * seq_len, hidden_size), dtype:[torch.half] + ln_weight: (hidden_size), dtype:[torch.half] + ln_bias:(hidden_size),dtype:[torch.half] + normalized_shape: list[int], [hidden_size] + Return: + output: (batch_count * seq_len, hidden_size), dtype:[torch.half] + + """ + return LayerNormFunction.apply( + input, ln_weight, ln_bias, output, normalized_shape, training + ) diff --git a/ixformer_sdk/train/functions/linear.py b/ixformer_sdk/train/functions/linear.py new file mode 100644 index 00000000..bcd6ca9a --- /dev/null +++ b/ixformer_sdk/train/functions/linear.py @@ -0,0 +1,89 @@ +import os +from typing import Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = ["linear"] + + +class LinearFunction(Function): + @staticmethod + def forward( + ctx, + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + output: torch.Tensor = None, + ): + if bias is not None: + if output is None: + output = ops.train.linear_forward(input, weight, bias) + else: + ops.train.linear_forward_(input, weight, bias, output) + else: + if output is None: + output = ops.train.linear_forward(input, weight) + else: + ops.train.linear_forward_(input, weight, output) + + ctx.has_bias = bias is not None + ctx.save_for_backward(input, weight) + + return output + + @staticmethod + def backward(ctx: FunctionCtx, dy: torch.Tensor): + x, w = ctx.saved_tensors + + dx = ops.train.linear_backward_dx(w, dy, x.shape) + + dw = ops.train.linear_backward_dw(x, dy, w.shape) + + if ctx.has_bias: + reduce_dims = list(range(dy.ndim - 1)) + db = torch.sum(dy, reduce_dims) + return dx, dw, db, None + else: + return dx, dw, None, None + + +def gemv_conditions(input, weight, bias, gemv_max_batch): + # gemv 使用的条件 input:[m,k] weight:[n,k] + # 1. m<=gemv_max_batch + # 2. k%2==0 n%2==0 + # 3. bias is None + input = input.view(-1, input.shape[-1]) + weight = weight.view(-1, weight.shape[-1]) + m = input.shape[0] + k = input.shape[1] + n = weight.shape[0] + if bias is None and m <= gemv_max_batch and k % 2 == 0 and n % 2 == 0: + return True + return False + + +def linear( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor = None, + output: torch.Tensor = None, + use_gemv: bool = True, + gemv_max_batch=1, +): + """ + Arguments: + input : [...,k] dtype: [torch.half, torch.bfloat16] + weights : [n,k] dtype: [torch.half, torch.bfloat16] + use_gemv: bool 是否使用gemv + gemv 使用的条件 input:[m,k] weight:[n,k] + 1. m<=gemv_max_batch + 2. k%2==0 n%2==0 + 3. bias is None + gemv_max_batch: int 用于是否满足gemv使用条件的判断 + Return: + output : [...,n] dtype: [torch.half, torch.bfloat16] + + """ + return LinearFunction.apply(input, weight, bias, output) diff --git a/ixformer_sdk/train/functions/matmul.py b/ixformer_sdk/train/functions/matmul.py new file mode 100644 index 00000000..01f813c3 --- /dev/null +++ b/ixformer_sdk/train/functions/matmul.py @@ -0,0 +1,108 @@ +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = ["matmul", "MatmulFunction"] + + +class MatmulFunction(Function): + @staticmethod + def forward( + ctx: FunctionCtx, + input: torch.Tensor, + other: torch.Tensor, + out: torch.Tensor = None, + transa: bool = False, + transb: bool = False, + alpha: float = 1.0, + beta: float = 0.0, + ): + ctx.save_for_backward(input, other) + ctx.params = (transa, transb, alpha, beta) + if out is None: + return ops.train.matmul( + input, other, transa=transa, transb=transb, alpha=alpha, beta=beta + ) + else: + return ops.train.matmul( + input, + other, + out=out, + transa=transa, + transb=transb, + alpha=alpha, + beta=beta, + ) + + @staticmethod + def backward(ctx: FunctionCtx, dy): + input, other = ctx.saved_tensors + transa, transb, alpha, beta = ctx.params + + if beta in [1, None]: + raise RuntimeError("Backward don't support beta == 1.0f") + + if not transa and not transb: + dx = matmul(dy, other, transb=True, alpha=alpha) + do = matmul(input, dy, transa=True, alpha=alpha) + return dx, do, None, None, None, None, None + + if transa and not transb: + dx = matmul(other, dy, transb=True, alpha=alpha) + do = matmul(input, dy, alpha=alpha) + return dx, do, None, None, None, None, None + + if not transa and transb: + dx = matmul(dy, other, alpha=alpha) + do = matmul(dy, input, transa=True, alpha=alpha) + return dx, do, None, None, None, None, None + + if transa and transb: + dx = matmul(other, dy, transa=True, transb=True, alpha=alpha) + do = matmul(dy, input, transa=True, transb=True, alpha=alpha) + return dx, do, None, None, None, None, None + + +def matmul( + input: torch.Tensor, + other: torch.Tensor, + *, + out: torch.Tensor = None, + transa: bool = False, + transb: bool = False, + alpha: float = 1.0, + beta: float = 0.0 +) -> torch.Tensor: + """ + 等价实现: + def pt_matmul(a, b, transa, transb, alpha): + if transa: + dims = list(range(a.ndim)) + dims[-1], dims[-2] = dims[-2], dims[-1] + a = a.permute(*dims).contiguous() + + if transb: + dims = list(range(b.ndim)) + dims[-1], dims[-2] = dims[-2], dims[-1] + b = b.permute(*dims).contiguous() + + return alpha * torch.matmul(a, b) + Arguments: + input: + 当transa为False shape : [...,m,k] dtype: torch.half + 当transa为True shape : [...,k,m] dtype: torch.half + other: + 当transb为False shape : [...,k,n] dtype: torch.half + 当transb为True shape : [...,n,k] dtype: torch.half + Return: + output: [...m,n] dtype: [torch.half] + + """ + if not input.is_contiguous(): + input = input.contiguous() + + if not other.is_contiguous(): + if not other.transpose(-2, -1).is_contiguous(): + other = other.contiguous() + + return MatmulFunction.apply(input, other, out, transa, transb, alpha, beta) diff --git a/ixformer_sdk/train/functions/residual_bias.py b/ixformer_sdk/train/functions/residual_bias.py new file mode 100644 index 00000000..8ccb4e32 --- /dev/null +++ b/ixformer_sdk/train/functions/residual_bias.py @@ -0,0 +1,82 @@ +from typing import Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +import ixformer + +__all__ = ["residual_bias"] + + +class ResidualBiasFunction(Function): + @staticmethod + def forward( + ctx, + input: torch.Tensor, + residual: torch.Tensor, + bias: torch.Tensor, + output: torch.Tensor, + alpha=1, + ): + if output is None: + output = torch.empty_like(input) + if alpha is None: + alpha = 1 + if bias is not None: + ops.train.add_residual_bias_forward(input, residual, bias, alpha, output) + else: + ops.train.add_residual_bias_forward(input, residual, alpha, output) + ctx.has_bias = bias is not None + ctx.alpha = alpha + return output + + @staticmethod + def backward(ctx: FunctionCtx, grad_output): + grad_input = torch.empty_like(grad_output) + grad_residual = torch.empty_like(grad_output) + if ctx.has_bias: + grad_bias = torch.empty( + [grad_output.size(-1)], + dtype=grad_output.dtype, + device=grad_output.device, + ) + ops.train.add_residual_bias_backward( + grad_output, + grad_input, + grad_residual, + grad_bias, + ctx.alpha, + ) + return (grad_input, grad_residual, grad_bias, None, None) + else: + ops.train.add_residual_bias_backward( + grad_output, + grad_input, + grad_residual, + ctx.alpha, + ) + return (grad_input, grad_residual, None, None, None) + + +def residual_bias( + input: torch.Tensor, + residual: torch.Tensor, + bias: torch.Tensor = None, + output: torch.Tensor = None, + alpha=1, +): + """ + 等价实现： + input = residual.float() * alpha + input.float() + bias.float() + + 参数说明: + Args: + input: shape:[batch_count, seq_len, hidden_size],dtype:[torch.half] + residual: shape:[batch_count, seq_len, hidden_size],dtype:[torch.half] + bias: shape:[hidden_size],dtype:[torch.half] + alpha: float + return: + output: shape:[batch_count, seq_len, hidden_size],dtype:[torch.half] + """ + return ResidualBiasFunction.apply(input, residual, bias, output, alpha) diff --git a/ixformer_sdk/train/functions/residual_bias_ln.py b/ixformer_sdk/train/functions/residual_bias_ln.py new file mode 100644 index 00000000..e45943ea --- /dev/null +++ b/ixformer_sdk/train/functions/residual_bias_ln.py @@ -0,0 +1,170 @@ +from typing import Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = ["residual_bias_ln"] + + +class ResidualBiasLnFunction(Function): + @staticmethod + def forward( + ctx, + input: torch.Tensor, + residual: torch.Tensor, + bias: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + output: torch.Tensor, + alpha=1, + is_post_ln=True, + ): + norm_size = ln_weight.size(-1) + + mean_size = input.numel() // norm_size + input_hat = torch.empty_like(input) + rstd = torch.empty([mean_size], dtype=input.dtype, device=input.device) + if bias is not None: + ops.train.add_residual_bias_ln_training_forward( + input, + residual, + bias, + ln_weight, + ln_bias, + alpha, + is_post_ln, + output, + input_hat, + rstd, + ) + else: + ops.train.add_residual_bias_ln_training_forward( + input, + residual, + ln_weight, + ln_bias, + alpha, + is_post_ln, + output, + input_hat, + rstd, + ) + ctx.norm_size = norm_size + ctx.has_bias = bias is not None + ctx.alpha = alpha + ctx.save_for_backward(input_hat, rstd, ln_weight) + + return output + + @staticmethod + def backward(ctx: FunctionCtx, grad_output): + input_hat, rstd_data, ln_weight = ctx.saved_tensors + + grad_input = torch.empty_like(input_hat) + grad_residual = torch.empty_like(input_hat) + grad_ln_weight = torch.empty_like(ln_weight) + grad_ln_bias = torch.empty_like(ln_weight) + + if ctx.has_bias: + grad_bias = torch.empty_like(ln_weight) + ops.train.add_residual_bias_ln_backward( + input_hat, + rstd_data, + ln_weight, + grad_output, + grad_ln_weight, + grad_ln_bias, + grad_input, + grad_residual, + grad_bias, + ctx.alpha, + ) + return ( + grad_input, + grad_residual, + grad_bias, + grad_ln_weight, + grad_ln_bias, + None, + None, + None, + None, + ) + else: + ops.train.add_residual_bias_ln_backward( + input_hat, + rstd_data, + ln_weight, + grad_output, + grad_ln_weight, + grad_ln_bias, + grad_input, + grad_residual, + ctx.alpha, + ) + return ( + grad_input, + grad_residual, + None, + grad_ln_weight, + grad_ln_bias, + None, + None, + None, + None, + ) + + +def residual_bias_ln( + input: torch.Tensor, + residual: torch.Tensor, + bias: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + alpha=1, + is_post_ln=True, + output: torch.Tensor = None, + training: bool = False, +): + """ + 等价实现： + input = residual.float() * alpha + input.float() + bias.float() + output = torch.nn.functional.layer_norm( + input, [input.shape[-1]], ln_weight.float(), ln_bias.float(), eps=1e-5) + + 参数说明: + Args: + input: shape:[batch_count * seq_len, hidden_size],dtype:[torch.half] + residual: shape:[batch_count * seq_len, hidden_size],dtype:[torch.half] + bias: shape:[hidden_size],dtype:[torch.half] + ln_weight:shape:[hidden_size],,dtype:[torch.half] + ln_bias:shape:[hidden_size],,dtype:[torch.half] + alpha: float + is_post_ln: bool, 是否应用layernorm 后处理 + return: + output: shape:[batch_count * seq_len, hidden_size],dtype:[torch.half] + """ + if alpha is None: + alpha = 1 + if not is_post_ln: + raise NotImplementedError() + if ln_weight is None or ln_bias is None: + raise NotImplementedError() + + if output is None: + output = torch.empty_like(input) + if not training: + if bias is not None: + ops.infer.add_residual_bias_ln_forward( + input, residual, bias, ln_weight, ln_bias, alpha, is_post_ln, output + ) + else: + ops.infer.add_residual_bias_ln_forward( + input, residual, ln_weight, ln_bias, alpha, is_post_ln, output + ) + return output + else: + return ResidualBiasLnFunction.apply( + input, residual, bias, ln_weight, ln_bias, output, alpha, is_post_ln + ) diff --git a/ixformer_sdk/train/functions/rms_norm.py b/ixformer_sdk/train/functions/rms_norm.py new file mode 100644 index 00000000..e079b4b7 --- /dev/null +++ b/ixformer_sdk/train/functions/rms_norm.py @@ -0,0 +1,324 @@ +import numbers +from typing import Union + +import ixformer._C as ops +import torch +from torch.nn import init +from torch.nn.parameter import Parameter + + +# apex interface for trainning add by xuelu.peng 2024/04/07 +class FusedRMSNormAffineFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, weight, normalized_shape, eps, memory_efficient=False, gradient_accumulation_fusion=False): + ctx.normalized_shape = normalized_shape + ctx.eps = eps + ctx.memory_efficient = memory_efficient + ctx.gradient_accumulation_fusion = gradient_accumulation_fusion + + input_ = input.contiguous() + weight_ = weight.contiguous() + output = torch.empty_like(input_) + normalized_shape_size = len(normalized_shape) + assert normalized_shape_size == 1 # 目前只支持normalized_shape_size=1 + invvar = torch.empty( + input_.shape[:-normalized_shape_size], + dtype=torch.float, + device=input_.device, + ) + ops.train.rms_norm_forward_training(input_, weight_, output, invvar, ctx.eps) + + ctx.save_for_backward(input_, weight_, invvar) + return output + + @staticmethod + def backward(ctx, grad_output): + input_, weight_, invvar = ctx.saved_tensors + + if ctx.gradient_accumulation_fusion: + if weight_.grad == None: + weight_.grad = torch.zeros_like(weight_) + grad_weight = weight_.grad + else: + grad_weight = torch.zeros_like(weight_) # 支持权重梯度累积融合，使用zeros_like，而不是emtpy_like 。 + + grad_input = torch.empty_like(input_) + + if input_.numel() < 4096 * 8192: + ops.train.rms_norm_backward_training( + input_, invvar, weight_, grad_output, grad_weight, grad_input + ) + else: ##llama 34b + ops.train.rms_norm_backward_training_opt( + input_, invvar, weight_, grad_output, grad_weight, grad_input + ) + + if ctx.gradient_accumulation_fusion: + grad_weight = None + return grad_input, grad_weight, None, None, None, None +def fused_rms_norm_affine( + input, weight, normalized_shape, eps=1e-6, memory_efficient=False, gradient_accumulation_fusion = False +): + return FusedRMSNormAffineFunction.apply( + input, weight, normalized_shape, eps, memory_efficient, gradient_accumulation_fusion + ) + + +class FusedRMSNorm(torch.nn.Module): + r"""Applies RMS Normalization over a mini-batch of inputs + + Currently only runs on cuda() tensors. + + .. math:: + y = \frac{x}{\mathrm{RMS}[x]} * \gamma + + The root-mean-square is calculated separately over the last + certain number dimensions which have to be of the shape specified by + :attr:`normalized_shape`. + :math:`\gamma` is a learnable affine transform parameter of + :attr:`normalized_shape` if :attr:`elementwise_affine` is ``True``. + `epsilon` is added to the mean-square, then the root of the sum is taken. + + .. note:: + Unlike Batch Normalization and Instance Normalization, which applies + scalar scale and bias for each entire channel/plane with the + :attr:`affine` option, RMS Normalization applies per-element scale + with :attr:`elementwise_affine`. + + This layer uses statistics computed from input data in both training and + evaluation modes. + + Args: + normalized_shape (int or list or torch.Size): input shape from an expected input + of size + + .. math:: + [* \times \text{normalized}\_\text{shape}[0] \times \text{normalized}\_\text{shape}[1] + \times \ldots \times \text{normalized}\_\text{shape}[-1]] + + If a single integer is used, it is treated as a singleton list, and this module will + normalize over the last dimension which is expected to be of that specific size. + eps: a value added to the denominator for numerical stability. Default: 1e-5 + elementwise_affine: a boolean value that when set to ``True``, this module + has learnable per-element affine parameters initialized to ones (for weights) + and zeros (for biases). Default: ``True``. + + Shape: + - Input: :math:`(N, *)` + - Output: :math:`(N, *)` (same shape as input) + + Examples:: + + >>> input = torch.randn(20, 5, 10, 10) + >>> # With Learnable Parameters + >>> m = ixformer.FusedRMSNorm(10) + >>> # Without Learnable Parameters + >>> m = ixformer.FusedRMSNorm(input.size()[1:], elementwise_affine=False) + >>> # Normalize over last dimension of size 10 #目前只支持在最后一维norm + >>> m = ixformer.FusedRMSNorm(10) + >>> # Activating the module + >>> output = m(input) + + .. _`Root Mean Square Layer Normalization`: https://arxiv.org/pdf/1910.07467.pdf + """ + + def __init__( + self, + normalized_shape, + eps=1e-5, + elementwise_affine=True, + memory_efficient=False, + gradient_accumulation_fusion=False + ): + super().__init__() + + if isinstance(normalized_shape, numbers.Integral): + normalized_shape = (normalized_shape,) + self.normalized_shape = torch.Size(normalized_shape) + self.eps = eps + self.elementwise_affine = elementwise_affine + self.memory_efficient = memory_efficient + self.gradient_accumulation_fusion = gradient_accumulation_fusion + + if self.elementwise_affine: + self.weight = Parameter(torch.empty(*normalized_shape)) + else: + self.register_parameter("weight", None) + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + init.ones_(self.weight) + + def forward(self, input): + if torch.jit.is_tracing() or torch.jit.is_scripting() or not input.is_cuda: + raise NotImplementedError() + + if self.elementwise_affine: + return fused_rms_norm_affine( + input, + self.weight, + self.normalized_shape, + self.eps, + self.memory_efficient, + self.gradient_accumulation_fusion + ) + else: + raise NotImplementedError() + + def extra_repr(self): + return "{normalized_shape}, eps={eps}, " "elementwise_affine={elementwise_affine}".format(**self.__dict__) + +class FusedRMSNormResFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, weight, residual, normalized_shape, eps, gradient_accumulation_fusion=False, memory_efficient=False): + ctx.normalized_shape = normalized_shape + ctx.eps = eps + ctx.memory_efficient = memory_efficient + ctx.gradient_accumulation_fusion = gradient_accumulation_fusion + + input_ = input.contiguous() + weight_ = weight.contiguous() + output = torch.empty_like(input_) + normalized_shape_size=len(normalized_shape) + assert normalized_shape_size == 1 #目前只支持normalized_shape_size=1 + invvar = torch.empty(input_.shape[:-normalized_shape_size], dtype=torch.float, device=input_.device) + + if residual is not None: + ctx.input_res = True + out_res = torch.empty_like(input_) + ops.train.rms_norm_res_forward_training(input_, weight_, output, invvar, ctx.eps, residual, out_res) + else: + ctx.input_res = False + ops.train.rms_norm_forward_training(input_, weight_, output, invvar, ctx.eps) + out_res = input_ + + # input_res 为 True 时 LN 的 input 为 input+redidual + ctx.save_for_backward(out_res, weight_, invvar) + return output, out_res + + @staticmethod + def backward(ctx, grad_output, grad_out_res): + input_, weight_, invvar = ctx.saved_tensors + + if ctx.gradient_accumulation_fusion: + if weight_.grad == None: + weight_.grad = torch.zeros_like(weight_) + grad_weight = weight_.grad + else: + grad_weight = torch.zeros_like(weight_) # 算子kernel 支持权重梯度累积融合，使用zeros_like，而不是emtpy_like 。 + + grad_input = torch.empty_like(input_) + + # rms_norm_res_backward_training 本身支持权重梯度累积融合，当不进行融合时，其输入 grad_weight 必须为 zero_like 。 + if input_.numel()< 4096*8192: + ops.train.rms_norm_res_backward_training(input_, invvar, weight_, + grad_output, grad_weight, grad_input, grad_out_res) + else:##llama 34b + ops.train.rms_norm_res_backward_training_opt(input_,invvar, weight_, + grad_output,grad_weight,grad_input,grad_out_res) + + if ctx.input_res: + grad_res = grad_input + else: + grad_res = None + + if ctx.gradient_accumulation_fusion: + grad_weight = None + + return grad_input, grad_weight, grad_res, None, None, None, None + +class FusedRMSNormRes(torch.nn.Module): + r"""Applies RMS Normalization and resdiual over a mini-batch of inputs, RMS Normalization part comes from FusedRMSNorm. + + Currently only runs on cuda() tensors. + + .. math:: + y = \frac{x}{\mathrm{RMS}[x]} * \gamma + + if residual None, x is input and output is equal to x, otherwise, x is input+residual and out_res is equal to x. + + The root-mean-square is calculated separately over the last + certain number dimensions which have to be of the shape specified by + :attr:`normalized_shape`. + :math:`\gamma` is a learnable affine transform parameter of + :attr:`normalized_shape` if :attr:`elementwise_affine` is ``True``. + `epsilon` is added to the mean-square, then the root of the sum is taken. + + .. note:: + Unlike Batch Normalization and Instance Normalization, which applies + scalar scale and bias for each entire channel/plane with the + :attr:`affine` option, RMS Normalization applies per-element scale + with :attr:`elementwise_affine`. + + This layer uses statistics computed from input data in both training and + evaluation modes. + + Args: + normalized_shape (int or list or torch.Size): input shape from an expected input + of size + + .. math:: + [* \times \text{normalized}\_\text{shape}[0] \times \text{normalized}\_\text{shape}[1] + \times \ldots \times \text{normalized}\_\text{shape}[-1]] + + If a single integer is used, it is treated as a singleton list, and this module will + normalize over the last dimension which is expected to be of that specific size. + eps: a value added to the denominator for numerical stability. Default: 1e-5 + elementwise_affine: a boolean value that when set to ``True``, this module + has learnable per-element affine parameters initialized to ones (for weights) + and zeros (for biases). Default: ``True``. + + Shape: + - Input: :math:`(N, *)` + - residual: :math:`(N, *)` (if not None) + - Output: :math:`(N, *)` (same shape as input) + - out_res: :math:`(N, *)` + + Examples:: + + >>> input = torch.randn(20, 5, 10, 10) + >>> res = torch.randn(20, 5, 10, 10) + >>> # With Learnable Parameters + >>> m = ixformer.FusedRMSNorm(10) + >>> # Without Learnable Parameters + >>> m = ixformer.FusedRMSNorm(input.size()[1:], elementwise_affine=False) + >>> # Normalize over last dimension of size 10 #目前只支持在最后一维norm + >>> m = ixformer.FusedRMSNorm(10) + >>> # Activating the module + >>> output, output_res = m(input, res) + + .. _`Root Mean Square Layer Normalization`: https://arxiv.org/pdf/1910.07467.pdf + """ + + def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True, memory_efficient=False, gradient_accumulation_fusion=False): + super().__init__() + + if isinstance(normalized_shape, numbers.Integral): + normalized_shape = (normalized_shape,) + self.normalized_shape = torch.Size(normalized_shape) + self.eps = eps + self.elementwise_affine = elementwise_affine + self.gradient_accumulation_fusion = gradient_accumulation_fusion + self.memory_efficient = memory_efficient + if self.elementwise_affine: + self.weight = Parameter(torch.empty(*normalized_shape)) + else: + self.register_parameter("weight", None) + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + init.ones_(self.weight) + + def forward(self, input, residual=None): + if torch.jit.is_tracing() or torch.jit.is_scripting() or not input.is_cuda: + raise NotImplementedError() + + if self.elementwise_affine: + return FusedRMSNormResFunction.apply(input, self.weight, residual, self.normalized_shape, self.eps, self.gradient_accumulation_fusion, self.memory_efficient) + else: + raise NotImplementedError() + + def extra_repr(self): + return "{normalized_shape}, eps={eps}, " "elementwise_affine={elementwise_affine}".format(**self.__dict__) diff --git a/ixformer_sdk/train/functions/swiglu.py b/ixformer_sdk/train/functions/swiglu.py new file mode 100644 index 00000000..dc657a0b --- /dev/null +++ b/ixformer_sdk/train/functions/swiglu.py @@ -0,0 +1,45 @@ +from typing import Union + +import ixformer._C as ops +import torch +from torch.autograd.function import Function, FunctionCtx + +__all__ = ["swiglu"] + + +class SwigluFunction(Function): + @staticmethod + def forward(ctx, input): + output_shape = list(input.shape) + output_shape[-1] = output_shape[-1] // 2 + output = input.new_empty(output_shape) + ops.train.swiglu_training_forward(input, output) + ctx.save_for_backward(input) + return output + + @staticmethod + def backward(ctx: FunctionCtx, grad_output): + input = ctx.saved_tensors[0] + grad_input = torch.empty_like(input) + ops.train.swiglu_training_backward(input, grad_output, grad_input) + return grad_input + + +def swiglu(input): + """ + 等价实现： + def ref_silu_and_mul(x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + x1, x2 = x.chunk(chunks=2, dim=-1) + res = torch.nn.functional.silu(x1) * x2 + return res.to(dtype) + + + 参数说明: + Args: + input: dtype:torch.float, torch.half, torch.bfloat16 + return: + output: dtype:torch.float, torch.half, torch.bfloat16 + """ + return SwigluFunction.apply(input) diff --git a/ixformer_sdk/train/speedformer/__init__.py b/ixformer_sdk/train/speedformer/__init__.py new file mode 100644 index 00000000..01eb1ffd --- /dev/null +++ b/ixformer_sdk/train/speedformer/__init__.py @@ -0,0 +1 @@ +from .speedformer import SpeedFormer \ No newline at end of file diff --git a/ixformer_sdk/train/speedformer/layers/__init__.py b/ixformer_sdk/train/speedformer/layers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/layers/baichuan/__init__.py b/ixformer_sdk/train/speedformer/layers/baichuan/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/layers/baichuan/attention.py b/ixformer_sdk/train/speedformer/layers/baichuan/attention.py new file mode 100644 index 00000000..ec7c402a --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/baichuan/attention.py @@ -0,0 +1,162 @@ +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from flash_attn import flash_attn_func, flash_attn_varlen_func +from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input +from ixformer.train.speedformer.models.baichuan.configuration_baichuan import BaichuanConfig +from ixformer.train.speedformer.models.baichuan.modeling_baichuan import Attention + +from ixformer.train.functions.fused_rope import fused_apply_rotary_pos_emb +from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding + +from ixformer.train.speedformer.layers.lazy import LazyInitContext + + +class FlashAttention(Attention): + # 这个类主要的改进包含：1. apply_rotary_pos_emb；2. flash-attn 代替 native attention + def __init__(self, config: BaichuanConfig): + super().__init__(config) + self.rotary_emb = RotaryEmbedding(self.head_dim) + + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: bool = False, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + proj = self.W_pack(hidden_states) + proj = proj.unflatten(-1, (3, self.hidden_size)).unsqueeze(0).transpose(0, -2).squeeze(-2) + + # fused_apply_rotary_pos_emb need qk to be in "sbhd", v stay in "bshd" + query_states = proj[0].view(bsz, q_len, self.num_heads, self.head_dim).transpose(0, 1).contiguous() + key_states = proj[1].view(bsz, q_len, self.num_heads, self.head_dim).transpose(0, 1).contiguous() + value_states = proj[2].view(bsz, q_len, self.num_heads, self.head_dim) + + kv_seq_len = key_states.shape[0] + if past_key_value is not None: + kv_seq_len += past_key_value[0].shape[0] + + # fused_apply_rotary_pos_emb need emb in float32 + emb = self.rotary_emb(kv_seq_len).to(dtype=torch.float32) + query_states = fused_apply_rotary_pos_emb(query_states, emb) + key_states = fused_apply_rotary_pos_emb(key_states, emb) + + if past_key_value is not None: + # reuse k, v, self_attention + key_states = torch.cat([past_key_value[0], key_states], dim=0) + value_states = torch.cat([past_key_value[1], value_states], dim=0) + + past_key_value = (key_states, value_states) if use_cache else None + + # after fused_apply_rotary_pos_emb, qk change to "bshd" for flashattn or "bhsd" for sdpa + if attention_mask is None: # flash-attn + query_states = query_states.transpose(0, 1).contiguous() + key_states = key_states.transpose(0, 1).contiguous() + else: # sdpa + query_states = query_states.permute(1, 2, 0, 3).contiguous() + key_states = key_states.permute(1, 2, 0, 3).contiguous() + value_states = value_states.transpose(1, 2).contiguous() + + ''' + if attention_mask is not None: + batch_size = query_states.shape[0] # bsz, q_len, self.num_heads, self.head_dim + query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( + query_states, key_states, value_states, attention_mask, q_len + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=0.0, + softmax_scale=None, + causal=True, + ) + + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, q_len) + else: + attn_output = flash_attn_func( + query_states, key_states, value_states, 0.0, softmax_scale=None, causal=True + ) + ''' + attn_output = self._flash_attention_forward( + query_states, key_states, value_states, q_len, attention_mask, dropout=0.0 + ) + + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def _flash_attention_forward( + self, + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + query_length: int, + attention_mask: Optional[torch.Tensor] = None, + dropout=0.0, + softmax_scale=None + ): + if attention_mask is not None: + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attention_mask, + dropout_p=0.0, + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal=query_length > 1, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + + else: + attn_output = flash_attn_func( + query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=self.is_causal + ) + + return attn_output + +class BaichuanAttention(FlashAttention): + def __init__(self) -> None: + raise NotImplementedError( + "BaichuanAttention is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native BaichuanAttention module to LlamaAttention module provided above." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + + LazyInitContext.materialize(module) + + # try to get normalized_shape, eps, elementwise_affine from the module + config = getattr(module, "config") + + attention = FlashAttention( + config=config, + ) + + attention.W_pack.weight = module.W_pack.weight + attention.o_proj.weight = module.o_proj.weight + + return attention diff --git a/ixformer_sdk/train/speedformer/layers/baichuan/baichuan_model.py b/ixformer_sdk/train/speedformer/layers/baichuan/baichuan_model.py new file mode 100644 index 00000000..0e057c77 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/baichuan/baichuan_model.py @@ -0,0 +1,141 @@ +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ixformer.train.speedformer.models.baichuan.configuration_baichuan import BaichuanConfig +from ixformer.train.speedformer.models.baichuan.modeling_baichuan import BaichuanModel +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.utils import logging, ContextManagers + + +from ixformer.train.speedformer.layers.lazy import LazyInitContext + +logger = logging.get_logger(__name__) + + +class IXFBaichuanModel(BaichuanModel): + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError( + "You have to specify either decoder_input_ids or decoder_inputs_embeds") + + seq_length_with_past = seq_length + past_key_values_length = 0 + + if past_key_values is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = () if use_cache else None + + for idx, decoder_layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + past_key_value = past_key_values[idx] if past_key_values is not None else None + + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + # None for past_key_value + return module(*inputs, output_attentions, None) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(decoder_layer), + hidden_states, + attention_mask, + position_ids, + None, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache += ( + layer_outputs[2 if output_attentions else 1],) + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) diff --git a/ixformer_sdk/train/speedformer/layers/baichuan/mlp.py b/ixformer_sdk/train/speedformer/layers/baichuan/mlp.py new file mode 100644 index 00000000..199cc6d0 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/baichuan/mlp.py @@ -0,0 +1,53 @@ +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +import ixformer.train.functions as F +from ixformer.train.speedformer.models.baichuan.configuration_baichuan import BaichuanConfig +from ixformer.train.speedformer.models.baichuan.modeling_baichuan import MLP +from transformers.utils import logging + +from ixformer.train.speedformer.layers.lazy import LazyInitContext + + +class BaseMLP(MLP): + """ + 这个层主要的优化点是：将linear1(act(cat(linear2(x), linear3(x))))的结构变成 linear1(act(linear23(x))) + """ + + def __init__(self, hidden_size, intermediate_size, hidden_act): + super().__init__(hidden_size, intermediate_size, hidden_act) + self.gate_up = nn.Linear( + hidden_size, intermediate_size * 2, bias=False) + del self.gate_proj, self.up_proj + del self.act_fn + + def forward(self, x): + res = self.gate_up(x) + down_proj = self.down_proj(F.swiglu(res)) + return down_proj + + +class IXFBaichuanMLP(BaseMLP): + def __init__(self) -> None: + raise NotImplementedError( + "IXFLlamaMLP is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to IXFLlamaMLP module provided above." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + hidden_size, intermediate_size = module.gate_proj.in_features, module.gate_proj.out_features + hidden_act = "silu" + + mlp = BaseMLP(hidden_size=hidden_size, + intermediate_size=intermediate_size, hidden_act=hidden_act) + + mlp.gate_up.weight.data = torch.concat( + (module.gate_proj.weight.data, module.up_proj.weight.data), dim=0) + mlp.down_proj.weight.data = module.down_proj.weight.data + + return mlp diff --git a/ixformer_sdk/train/speedformer/layers/bloom/__init__.py b/ixformer_sdk/train/speedformer/layers/bloom/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/layers/bloom/attention.py b/ixformer_sdk/train/speedformer/layers/bloom/attention.py new file mode 100644 index 00000000..e9e247ad --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/bloom/attention.py @@ -0,0 +1,160 @@ +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from flash_attn import flash_attn_func, flash_attn_varlen_func +from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input +from ixformer.train.speedformer.models.bloom.modeling_bloom import BloomAttention, dropout_add +from ixformer.train.speedformer.models.bloom.configuration_bloom import BloomConfig + +from apex.transformer.functional.fused_rope import fused_apply_rotary_pos_emb_cached +from apex.transformer.functional.fused_rope import FusedRoPEFunc + + +class FlashAttention(BloomAttention): + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + alibi: torch.Tensor, + attention_mask: torch.Tensor, + layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + head_mask: Optional[torch.Tensor] = None, + use_cache: bool = False, + output_attentions: bool = False, + ): + fused_qkv = self.query_key_value(hidden_states) + (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv) # 3 x [batch_size, seq_length, num_heads, head_dim] + batch_size, q_length, _, _ = query_layer.shape + + if layer_past is not None: + past_key, past_value = layer_past + key_layer = torch.cat((past_key, key_layer), dim=1) + value_layer = torch.cat((past_value, value_layer), dim=1) + + present = (key_layer, value_layer) if use_cache else None + # if attention_mask is not None: + if False: + query_layer, key_layer, value_layer, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( + query_layer, key_layer, value_layer, attention_mask, q_length + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + attn_output_unpad = flash_attn_varlen_func( + query_layer, + key_layer, + value_layer, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=0.0, + softmax_scale=None, + causal=True, + use_alibi=True, + ) + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, q_length) + else: + attn_output = flash_attn_func( + query_layer, key_layer, value_layer, 0.0, softmax_scale=None, causal=True, use_alibi=True, + ) + + attn_output = attn_output.reshape(batch_size, q_length, attn_output.shape[2]*attn_output.shape[3]).contiguous() + output_tensor = self.dense(attn_output) + + output_tensor = dropout_add(output_tensor, residual, self.hidden_dropout, self.training) + + outputs = (output_tensor, present, None) + + return outputs + + + def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): + + def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = index_first_axis( + key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +class BloomFlashAttention(FlashAttention): + + def __init__(self) -> None: + raise NotImplementedError( + "BloomAttention is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native BloomAttention module to FlashAttention module provided above." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + + # try to get normalized_shape, eps, elementwise_affine from the module + new_config = BloomConfig() + new_config.pretraining_tp = module.pretraining_tp + new_config.slow_but_exact = module.slow_but_exact + new_config.hidden_size = module.hidden_size + new_config.n_head = module.num_heads + new_config.hidden_size = module.split_size + new_config.hidden_dropout = module.hidden_dropout + new_config.attention_dropout = module.attention_dropout.p + + attention = FlashAttention( + config=new_config, + ) + + attention.query_key_value.weight = module.query_key_value.weight + attention.query_key_value.bias = module.query_key_value.bias + + attention.dense.weight = module.dense.weight + attention.dense.bias = module.dense.bias + + return attention \ No newline at end of file diff --git a/ixformer_sdk/train/speedformer/layers/chatglm/__init__.py b/ixformer_sdk/train/speedformer/layers/chatglm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/layers/chatglm/attention.py b/ixformer_sdk/train/speedformer/layers/chatglm/attention.py new file mode 100644 index 00000000..1ba35134 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/chatglm/attention.py @@ -0,0 +1,199 @@ +import math +import os +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ixformer.train.speedformer.models.chatglm.modeling_chatglm import ( + CoreAttention, + SelfAttention, + split_tensor_along_last_dim, + apply_rotary_pos_emb +) +from ixformer.train.speedformer.models.chatglm.configuration_chatglm import ChatGLMConfig + +from transformers.utils import is_flash_attn_2_available + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + +class FlashCoreAttention(CoreAttention): + + def forward(self, query_layer, key_layer, value_layer, attention_mask): + if int(os.environ.get("USE_FLASH_ATTN", 0)): + query_layer, key_layer, value_layer = [ + k.permute(1, 0, 2, 3) for k in [query_layer, key_layer, value_layer]] + batch_size, query_length, _, _ = query_layer.shape + + if attention_mask is not None: + batch_size = query_layer.shape[0] + query_layer, key_layer, value_layer, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( + query_layer, key_layer, value_layer, attention_mask, query_length + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + attn_output_unpad = flash_attn_varlen_func( + query_layer, + key_layer, + value_layer, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=0.0, + softmax_scale=None, + causal=True, + ) + attn_output = pad_input( + attn_output_unpad, indices_q, batch_size, query_length) + context_layer = attn_output.permute(1, 0, 2, 3) + else: + attn_output = flash_attn_func( + query_layer, key_layer, value_layer, 0.0, softmax_scale=None, causal=True + ) + context_layer = attn_output.permute(1, 0, 2, 3) + + if attention_mask is not None: + if query_layer.shape[2] != key_layer.shape[2]: + num_group = query_layer.shape[2] // key_layer.shape[2] + final_shape = (*key_layer.shape[:2], *query_layer.shape[2:]) + key_layer = key_layer.unsqueeze(-2) + key_layer = key_layer.expand( + -1, -1, -1, num_group, -1 + ) + key_layer = key_layer.contiguous().view( + final_shape + ) + value_layer = value_layer.unsqueeze(-2) + value_layer = value_layer.expand( + -1, -1, -1, num_group, -1 + ) + value_layer = value_layer.contiguous().view( + final_shape + ) + + query_layer, key_layer, value_layer = [ + k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]] # bhsd + attention_mask = ~attention_mask + context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, + attention_mask) + context_layer = context_layer.permute(2, 0, 1, 3) + + else: + query_layer, key_layer, value_layer = [ + k.permute(1, 0, 2, 3) for k in [query_layer, key_layer, value_layer]] # bshd + context_layer = flash_attn_func( + query_layer, key_layer, value_layer, 0, softmax_scale=None, causal=True + ) # bshd + context_layer = context_layer.permute(1, 0, 2, 3) + + context_layer = context_layer.reshape( + context_layer.size(0), context_layer.size(1), -1) + + return context_layer + + +class FlashSelfAttention(SelfAttention): + + def __init__(self, config: ChatGLMConfig, layer_number, device=None): + super().__init__(config, layer_number, device=device) + self.core_attention = FlashCoreAttention(config, self.layer_number) + + def forward(self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True): + mixed_x_layer = self.query_key_value(hidden_states) + if self.multi_query_attention: + (query_layer, key_layer, value_layer) = mixed_x_layer.split( + [ + self.num_attention_heads_per_partition * self.hidden_size_per_attention_head, + self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, + self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, + ], + dim=-1, + ) + query_layer = query_layer.view( + query_layer.size()[ + :-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + key_layer = key_layer.view( + key_layer.size()[ + :-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head) + ) + value_layer = value_layer.view( + value_layer.size()[:-1] + + (self.num_multi_query_groups_per_partition, + self.hidden_size_per_attention_head) + ) + else: + new_tensor_shape = mixed_x_layer.size()[:-1] + \ + (self.num_attention_heads_per_partition, + 3 * self.hidden_size_per_attention_head) + mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) + + # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn] + (query_layer, key_layer, value_layer) = split_tensor_along_last_dim( + mixed_x_layer, 3) + + if rotary_pos_emb is not None: + query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb) + key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb) + + # adjust key and value for inference + if kv_cache is not None: + cache_k, cache_v = kv_cache + key_layer = torch.cat((cache_k, key_layer), dim=0) + value_layer = torch.cat((cache_v, value_layer), dim=0) + if use_cache: + kv_cache = (key_layer, value_layer) + else: + kv_cache = None + + # 这里省略了 kv "sbhd" -> "sb(h*num_multi-group)d" 的过程，因为flash-attn支持 MGA + # ================================== + # core attention computation + # ================================== + + context_layer = self.core_attention( + query_layer, key_layer, value_layer, attention_mask) + + # ================= + # Output. [sq, b, h] + # ================= + + output = self.dense(context_layer) + + return output, kv_cache + + +class ChatglmFlashAttention(FlashSelfAttention): + + def __init__(self) -> None: + raise NotImplementedError( + "BloomAttention is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native BloomAttention module to FlashAttention module provided above." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + # 这个原实现没有在类中保存config，所以需要初始化一个config + layer_number = getattr(module, "layer_number") + config = getattr(module, "config") + attention = FlashSelfAttention( + config=config, + layer_number=layer_number, + ) + + attention.query_key_value.weight.data = module.query_key_value.weight.data + attention.dense.weight.data = module.dense.weight.data + if getattr(attention.query_key_value, "bias") is not None: + attention.query_key_value.bias.data = module.query_key_value.bias.data + if getattr(attention.dense, "bias") is not None: + attention.dense.bias.data = module.dense.bias.data + + return attention diff --git a/ixformer_sdk/train/speedformer/layers/chatglm/attributions.py b/ixformer_sdk/train/speedformer/layers/chatglm/attributions.py new file mode 100644 index 00000000..02b71408 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/chatglm/attributions.py @@ -0,0 +1,9 @@ +from ixformer.train.speedformer.models.chatglm.modeling_chatglm import RotaryEmbedding +from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding + + +class ChatglmRotaryEmbedding(RotaryEmbedding): + def from_native_attr(attr_class, *args, **kwargs): + dim = attr_class.dim + rote = RotaryEmbedding(dim=dim) + return rote diff --git a/ixformer_sdk/train/speedformer/layers/chatglm/methods.py b/ixformer_sdk/train/speedformer/layers/chatglm/methods.py new file mode 100644 index 00000000..4ed2a29d --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/chatglm/methods.py @@ -0,0 +1,127 @@ +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ixformer.train.speedformer.models.chatglm.modeling_chatglm import ChatGLMModel + + +def ChatGLMModel_forward(): + from transformers.modeling_outputs import BaseModelOutputWithPast + from transformers.utils import logging, is_flash_attn_2_available + + def forward( + self, + input_ids, + position_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.BoolTensor] = None, + full_attention_mask: Optional[torch.BoolTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + inputs_embeds: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + def is_lower_triangular(mask): + """ + ixdnn 虽然支持2种causal mask, 如下图: + mode0: + if seqlen_q < seqlen_k + 1 0 0 0 0 + 1 1 0 0 0 + if seqlen_k < seqlen_q + 1 0 + 1 1 + 1 1 + 1 1 + 1 1 + mode1: + if seqlen_q < seqlen_k + 1 1 1 1 0 + 1 1 1 1 1 + if seqlen_k < seqlen_q + 0 0 + 0 0 + 0 0 + 1 0 + 1 1 + + 但 flash-attn 目前只支持 mode1, 所以下面需要判断一下传入的mask是不是mode1这种模式 + """ + batch_size, _, rows, cols = mask.shape + + # 创建一个mode1的下三角矩阵 + if rows <= cols: + part = torch.ones(rows, cols - rows, + dtype=torch.bool, device=mask.device) + gt = ~torch.triu(torch.ones( + rows, rows, dtype=torch.bool, device=mask.device), diagonal=1) + gt = torch.cat((part, gt), dim=1) + else: + part = torch.zeros( + rows-cols, cols, dtype=torch.bool, device=mask.device) + gt = ~torch.triu(torch.ones( + cols, cols, dtype=torch.bool, device=mask.device), diagonal=1) + gt = torch.cat((part, gt), dim=0) + gt = gt[None, None, :, :].expand(batch_size, -1, -1, -1) + + # 检查所有的元素是不是都一样 + check = (gt == mask).all() + + return check + + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + batch_size, seq_length = input_ids.shape + + if inputs_embeds is None: + inputs_embeds = self.embedding(input_ids) + + if self.pre_seq_len is not None: + if past_key_values is None: + past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device, + dtype=inputs_embeds.dtype) + if attention_mask is not None: + attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)), + attention_mask], dim=-1) + + if full_attention_mask is None: + if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1): + full_attention_mask = self.get_masks( + input_ids, past_key_values, padding_mask=attention_mask) + + # Rotary positional embeddings + rotary_pos_emb = self.rotary_pos_emb(self.seq_length) + if position_ids is not None: + rotary_pos_emb = rotary_pos_emb[position_ids] + else: + rotary_pos_emb = rotary_pos_emb[None, :seq_length] + rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous() + + # Run encoder. + attn_mask = None + if full_attention_mask is not None: + if not is_lower_triangular(full_attention_mask): + attn_mask = full_attention_mask + hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder( + inputs_embeds, attn_mask, rotary_pos_emb=rotary_pos_emb, + kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states + ) + + if not return_dict: + return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=presents, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + return forward diff --git a/ixformer_sdk/train/speedformer/layers/cross_entropy_loss.py b/ixformer_sdk/train/speedformer/layers/cross_entropy_loss.py new file mode 100644 index 00000000..ce9c91a3 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/cross_entropy_loss.py @@ -0,0 +1,370 @@ +import time +import numpy as np +import torch +import triton +import triton.language as tl +from packaging.version import Version +if Version(triton.__version__) >= Version("3.0.0"): + from triton.language.extra import libdevice + triton_tanh = libdevice.tanh +else: + import triton.language as tl + triton_tanh = tl.math.tanh + + +def calculate_settings(n): + BLOCK_SIZE = triton.next_power_of_2(n) + if BLOCK_SIZE > MAX_FUSED_SIZE: + raise RuntimeError(f"Cannot launch Triton kernel since n = {n} exceeds " + f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.") + num_warps = 4 + if BLOCK_SIZE >= 32768: + num_warps = 32 + elif BLOCK_SIZE >= 8192: + num_warps = 16 + elif BLOCK_SIZE >= 2048: + num_warps = 8 + return BLOCK_SIZE, num_warps + + +@triton.heuristics({"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING"], }) +@triton.jit +def _cross_entropy_forward( + logits_ptr, logits_row_stride, + loss_ptr, + logsumexp_ptr, + labels_ptr, + VOCAB_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + DO_SOFTCAPPING: tl.constexpr, + SOFTCAP: tl.constexpr, +): + """ + Cross Entropy Loss = 1/n sum [ -yi log(Pi) ] + Pi = exp(xi) / sum(exp(xi)) + CE_i = -y log(p) = -y log[ exp(x) / sum(exp(x)) ] + = -y [ x - log[sum(exp(x))] ] + = y * (log[sum(exp(x))] - x) + If y == 0: CE_i = 0 + If y == 1: CE_i = logsumexp - x + + logsumexp is also stable + Take y = log[sum(exp(x))] + exp(y) = sum(exp(x)) + exp(y) = sum(exp(x - c)*exp(c)) Since e^(x-c)*e^c = e^x + exp(y) = exp(c)*sum(exp(x - c)) + y = log(exp(c)*sum(exp(x - c))) + y = c + log[sum(exp(x - c))] + This means we can set c = max(x) to make sure + exp(x - c) always is exp(x - max(x)). + This ensures exp(x - max(x))'s maximum is 1 as exp(0) = 1. + """ + row_idx = tl.program_id(0) + logits_ptr += row_idx * logits_row_stride.to(tl.int64) + loss_ptr += row_idx + logsumexp_ptr += row_idx + labels_ptr += row_idx + + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < VOCAB_SIZE + + label_idx = tl.load(labels_ptr).to(tl.int32) + logits = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf")) + # Do logit softcapping for Gemma 2: t * tanh(1/t * x) + if DO_SOFTCAPPING: + logits = SOFTCAP * triton_tanh(logits / SOFTCAP) + + logits = logits.to(tl.float32) + c = tl.max(logits, 0) + logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0)) + + if label_idx != -100: + x = tl.load(logits_ptr + label_idx) + # Do logit softcapping for Gemma 2: t * tanh(1/t * x) + if DO_SOFTCAPPING: + x = SOFTCAP * triton_tanh(x / SOFTCAP) + loss = logsumexp - x.to(tl.float32) + else: + loss = 0.0 + tl.store(logsumexp_ptr, logsumexp) + tl.store(loss_ptr, loss) + + +@triton.heuristics({"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING"], }) +@triton.jit +def _chunked_cross_entropy_forward( + logits_ptr, logits_row_stride, + loss_ptr, + logsumexp_ptr, + labels_ptr, + VOCAB_SIZE: tl.constexpr, + N_CHUNKS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + DO_SOFTCAPPING: tl.constexpr, + SOFTCAP: tl.constexpr, +): + """ + 256K vocab divided in 4 chunks + + |-65536-| |-65536-| |-65536-| |-65536-| + |-------| |-------| |-------| |-------| + |-------| |-------| |-------| |-------| + + If y == 0: CE_i = 0 + If y == 1: CE_i = logsumexp - x + + Notice we can do logsumexp for each chunk and then + logsumexp[chunk_sum(logsumexp)] == logsumexp + + chunk_sum = log[chunk_sum(logsumexp)] + = log[exp(logsumexp(a)) + ... + exp(logsumexp(z))] + = log[exp(log[sum(exp(a))]) + ... + exp(log[sum(exp(z))])] + = log[sum(exp(a)) + ... + sum(exp(z))] + = logsumexp(x) + + This means we can perform a logsumexp for each chunk, then do a + final logsumexp reduction! + + Ie do: logsumexp(chunked_logsumexp) - x + """ + row_idx = tl.program_id(0) + chunk_idx = tl.program_id(1) + logits_ptr += row_idx * logits_row_stride.to(tl.int64) + loss_ptr += row_idx + logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx + labels_ptr += row_idx + + col_offsets = chunk_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < VOCAB_SIZE + + label_idx = tl.load(labels_ptr).to(tl.int32) + logits = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf")) + # Do logit softcapping for Gemma 2: t * tanh(1/t * x) + if DO_SOFTCAPPING: + logits = SOFTCAP * triton_tanh(logits / SOFTCAP) + + logits = logits.to(tl.float32) + c = tl.max(logits, 0) + logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0)) + + if chunk_idx == 0: + # logsumexp(chunked_logsumexp) - x + # Do the -x separately + if label_idx != -100: + x = tl.load(logits_ptr + label_idx).to(tl.float32) + # Do logit softcapping for Gemma 2: t * tanh(1/t * x) + if DO_SOFTCAPPING: + x = SOFTCAP * triton_tanh(x / SOFTCAP) + loss = -1.0 * x.to(tl.float32) + else: + loss = 0.0 + tl.store(loss_ptr, loss) + + tl.store(logsumexp_ptr, logsumexp) + + +@triton.heuristics({"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING"], }) +@triton.jit +def _cross_entropy_backward( + logits_ptr, logits_row_stride, + dloss_ptr, dloss_row_stride, + logsumexp_ptr, + labels_ptr, + VOCAB_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + DO_SOFTCAPPING: tl.constexpr, + SOFTCAP: tl.constexpr, +): + """ + CE_i = -y log(P) = y * (log[sum(exp(x))] - x) + dC/dx = d/dx (y * log[sum(exp(x))] - x * y) + + From https://en.wikipedia.org/wiki/LogSumExp + d/dx logsumexp = exp(x) / sum(exp(x)) = softmax(x) + + dC/dx = y * exp(x) / sum(exp(x)) - d/dx (x * y) + dC/dx = y * exp[ log[exp(x) / sum(exp(x))] ] using x = exp(log(x)) trick + dC/dx = y * exp[x - logsumexp] - d/dx (x * y) + + If y == 0: dC/dx = 0 + If y == 1 and x == label: dC/dlabel = exp[x - logsumexp] - 1 + If y == 1 and x != label: dC/dx = exp[x - logsumexp] + """ + row_idx = tl.program_id(0) + block_idx = tl.program_id(1) + + logits_ptr += row_idx * logits_row_stride.to(tl.int64) + dloss_ptr += row_idx * dloss_row_stride + col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < VOCAB_SIZE + label_idx = tl.load(labels_ptr + row_idx).to(tl.int32) + + if label_idx != -100: + dloss = tl.load(dloss_ptr) + else: + dloss = 0.0 + + x = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf")) + # Do logit softcapping for Gemma 2: t * tanh(1/t * x) + if DO_SOFTCAPPING: + # d/dx [t * tanh(1/t * x)] = 1 - tanh^2(1/t * x) + partial = triton_tanh(x / SOFTCAP) + x = SOFTCAP * partial + + logsumexp = tl.load(logsumexp_ptr + row_idx) + y = tl.exp(x.to(tl.float32) - logsumexp) + y = tl.where( + col_offsets == label_idx, + y - 1.0, # exp(x - logsumexp) - 1 + y, # exp(x - logsumexp) + ) + + if DO_SOFTCAPPING: + # d/dx [t * tanh(1/t * x)] = 1 - tanh^2(1/t * x) + y = y * (1.0 - partial*partial) + + # If y == 0: dC/dx = 0 ==> we already masked it to be = 0, so dloss = 0. + tl.store(logits_ptr + col_offsets, dloss * y, mask=mask) + + +MAX_FUSED_SIZE = 65536 # 2**16 + + +class Fast_CrossEntropyLoss(torch.autograd.Function): + @staticmethod + def forward(ctx, logits, labels, logit_softcapping=0): + n_rows, vocab_size = logits.shape + + div, mod = divmod(vocab_size, MAX_FUSED_SIZE) + n_chunks = div + (mod != 0) + losses = torch.empty(n_rows, dtype=torch.float32, device=logits.device) + + DO_SOFTCAPPING = (logit_softcapping != 0) + + if n_chunks == 1: + # For small vocabs <= 65336 like Llama, Mistral + BLOCK_SIZE, num_warps = calculate_settings(vocab_size) + logsumexp = torch.empty( + n_rows, dtype=torch.float32, device=logits.device) + + _cross_entropy_forward[(n_rows,)]( + logits, logits.stride(0), + losses, + logsumexp, + labels, + VOCAB_SIZE=vocab_size, + BLOCK_SIZE=BLOCK_SIZE, + DO_SOFTCAPPING=DO_SOFTCAPPING, + SOFTCAP=logit_softcapping, + num_warps=num_warps, + ) + else: + # For large vocabs > 65336 like Gemma 256K + logsumexp = torch.empty( + (n_rows, n_chunks,), dtype=torch.float32, device=logits.device) + + _chunked_cross_entropy_forward[(n_rows, n_chunks,)]( + logits, logits.stride(0), + losses, + logsumexp, + labels, + VOCAB_SIZE=vocab_size, + N_CHUNKS=n_chunks, + BLOCK_SIZE=MAX_FUSED_SIZE, + DO_SOFTCAPPING=DO_SOFTCAPPING, + SOFTCAP=logit_softcapping, + num_warps=32, + ) + # logsumexp(chunked_logsumexp) - x + # Do the -x separately + logsumexp = torch.logsumexp(logsumexp, dim=1) # Row sum + losses += logsumexp + # Don't forget to mask padding out! + losses.masked_fill_(labels == -100, 0) + + ctx.save_for_backward(logits, logsumexp, labels) + ctx.DO_SOFTCAPPING = DO_SOFTCAPPING + ctx.logit_softcapping = logit_softcapping + return losses + + @staticmethod + def backward(ctx, dlosses): + logits, logsumexp, labels = ctx.saved_tensors + n_rows, vocab_size = logits.shape + + BLOCK_SIZE = 4096 + div, mod = divmod(vocab_size, BLOCK_SIZE) + n_blocks = div + (mod != 0) + + _cross_entropy_backward[(n_rows, n_blocks,)]( + logits, logits.stride(0), + dlosses, dlosses.stride(0), + logsumexp, + labels, + VOCAB_SIZE=vocab_size, + BLOCK_SIZE=BLOCK_SIZE, + DO_SOFTCAPPING=ctx.DO_SOFTCAPPING, + SOFTCAP=ctx.logit_softcapping, + num_warps=8, + ) + return logits, None, None, + + +@torch._disable_dynamo +def fast_cross_entropy_loss(logits, labels, logit_softcapping=0): + """ + Arguments: + logits: (batch, seq_len, vocab_size) + labels: (batch, seq_len,) + Returns: + losses: float + """ + assert len(logits.size()) == 2 or len(logits.size()) == 3 + if len(logits.size()) == 3: + batch, seq_len, d = logits.shape + assert (labels.shape == (batch, seq_len)) + logits = logits.view(batch*seq_len, d) + labels = labels.view(-1) + + loss = Fast_CrossEntropyLoss.apply( + logits, + labels, + logit_softcapping, + ) + n_items = torch.count_nonzero(labels != -100) + return loss.sum() / n_items + + +if __name__ == "__main__": + shift_logits_numpy = np.random.randn(4096, 32000).astype(np.float32) + shift_labels_numpy = np.random.randint(0, 32000, (4096, )).astype(np.int64) + + shift_logits = torch.from_numpy(shift_logits_numpy).cuda() + shift_labels = torch.from_numpy(shift_labels_numpy).cuda() + + shift_logits_ref = torch.from_numpy(shift_logits_numpy).cuda() + shift_labels_ref = torch.from_numpy(shift_labels_numpy).cuda() + + shift_logits.requires_grad = True + shift_logits_ref.requires_grad = True + + # test accuracy + loss = fast_cross_entropy_loss(shift_logits, shift_labels) + loss_ref = torch.nn.CrossEntropyLoss()(shift_logits_ref, shift_labels_ref) + loss.backward() + loss_ref.backward() + + torch.testing.assert_close(loss, loss_ref) + torch.testing.assert_close(shift_logits.grad, shift_logits_ref.grad) + + start = time.time() + for i in range(1000): + loss = fast_cross_entropy_loss(shift_logits, shift_labels) + loss.backward() + print("triton:", time.time() - start) + + start = time.time() + for i in range(1000): + loss_ref = torch.nn.CrossEntropyLoss()(shift_logits, shift_labels) + loss_ref.backward() + print("torch:", time.time() - start) diff --git a/ixformer_sdk/train/speedformer/layers/fast_lora/__init__.py b/ixformer_sdk/train/speedformer/layers/fast_lora/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/layers/fast_lora/fast_lora.py b/ixformer_sdk/train/speedformer/layers/fast_lora/fast_lora.py new file mode 100644 index 00000000..e8adb950 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/fast_lora/fast_lora.py @@ -0,0 +1,305 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# 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. + +from ixformer.train.speedformer.layers.fast_lora.swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel +import torch +from ixformer.train.speedformer.layers.fast_lora.utils import ( + fast_dequantize, + QUANT_STATE, + get_lora_parameters, + matmul_lora, + torch_amp_custom_fwd, + torch_amp_custom_bwd, +) + + +class LoRA_MLP(torch.autograd.Function): + """ + ### LoRA weights + G = G + Ag @ Bg + U = U + Au @ Bu + W = W + Aw @ Bw + + ### SwiGLU(X) + e = X @ G + f = e * sigmoid(e) + g = X @ U + h = f * g + i = h @ W + + ### Backpropagation chain rule + See our blog post for more details + + df = sigmoid(e) * (1 - f) + f + dC/dW = h.T @ dY + dC/dU = X.T @ (D @ W.T * f) + dC/dG = X.T @ (D @ W.T * df * g) + + ### Down projection LoRA weights + dC/dAw = dC/dW @ B.T + dC/dBw = A.T @ dC/dW + dC/dAw = h.T @ dY @ B.T + dC/dBw = A.T @ h.T @ dY + + ### Up projection LoRA weights + dC/dAu = X.T @ (D @ W.T * f) @ B.T + dC/dBu = A.T @ X.T @ (D @ W.T * f) + + ### Gate projection LoRA weights + dC/dAg = X.T @ (D @ W.T * df * g) @ B.T + dC/dBg = A.T @ X.T @ (D @ W.T * df * g) + + Don't forget to see our blog post for more details! + """ + @staticmethod + @torch_amp_custom_fwd + def forward(ctx, X: torch.Tensor, + gateW, gateW_quant, gateA, gateB, gateS, + upW, upW_quant, upA, upB, upS, + downW, downW_quant, downA, downB, downS, + _forward_function, _backward_function,): + dtype = X.dtype + + e = matmul_lora(X, gateW, gateW_quant, gateA, gateB, gateS) + g = matmul_lora(X, upW, upW_quant, upA, upB, upS) + h = _forward_function(e, g) + i = matmul_lora(h, downW, downW_quant, downA, downB, downS) + + ctx.custom_saved_tensors = ( + gateW, gateW_quant, gateS, + upW, upW_quant, upS, + downW, downW_quant, downS, + _backward_function, + ) + ctx.save_for_backward(gateA, gateB, upA, upB, downA, downB, + X, e, g) + return i + pass + + @staticmethod + @torch_amp_custom_bwd + def backward(ctx, dY: torch.Tensor): + gateW, gateW_quant, gateS, upW, upW_quant, upS, downW, downW_quant, downS, \ + _backward_function = ctx.custom_saved_tensors + gateA, gateB, upA, upB, downA, downB, \ + X, e, g = ctx.saved_tensors + + gateA, gateB, upA, upB, downA, downB = \ + gateA.t(), gateB.t(), upA.t(), upB.t(), downA.t(), downB.t() + + batch, seq_len, hd = X.shape + dY = dY.view(-1, dY.shape[-1]) + X = X .view(-1, X .shape[-1]) + e = e .view(-1, e .shape[-1]) + g = g .view(-1, g .shape[-1]) + dtype = X.dtype + + DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) + DW, e, g = _backward_function(DW, e, g) + h, df, de = DW, e, g + + # Down projection LoRA weights + d_downA = h.t() @ (dY @ downB.t()) + d_downB = (downA.t() @ h.t()) @ dY + d_downA *= downS + d_downB *= downS + + # Up projection LoRA weights + d_upA = X.t() @ (df @ upB.t()) + d_upB = (upA.t() @ X.t()) @ df + d_upA *= upS + d_upB *= upS + + # Gate projection LoRA weights + d_gateA = X.t() @ (de @ gateB.t()) + d_gateB = (gateA.t() @ X.t()) @ de + d_gateA *= gateS + d_gateB *= gateS + + # dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS) + # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS) + upW = fast_dequantize(upW.t(), upW_quant) + dX = torch.matmul(df, upW.t(), out=X) + del upW + dX += df @ upB.to(dtype).t() @ (upS * upA.to(dtype).t()) + + gateW = fast_dequantize(gateW.t(), gateW_quant) + dX += de @ gateW.t() + del gateW + dX += de @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t()) + + # gateW, gateW_quant, gateA, gateB, gateS, + # upW, upW_quant, upA, upB, upS, + # downW, downW_quant, downA, downB, downS, + return dX.view(batch, seq_len, hd), \ + None, None, d_gateA.t(), d_gateB.t(), None, \ + None, None, d_upA.t(), d_upB.t(), None, \ + None, None, d_downA.t(), d_downB.t(), None, \ + None, None, # _backward and _forward + pass + + +pass + + +def apply_lora_mlp_swiglu(self, X): + gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters( + self.gate_proj) + upW, upW_quant, upA, upB, upS = get_lora_parameters( + self. up_proj) + downW, downW_quant, downA, downB, downS = get_lora_parameters( + self.down_proj) + + out = LoRA_MLP.apply(X, + gateW, gateW_quant, gateA, gateB, gateS, + upW, upW_quant, upA, upB, upS, + downW, downW_quant, downA, downB, downS, + swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel,) + return out + + +pass + + +class LoRA_FUSEMLP(torch.autograd.Function): + """ + ### LoRA weights + G = G + Ag @ Bg + U = U + Au @ Bu + W = W + Aw @ Bw + + ### SwiGLU(X) + e = X @ G + f = e * sigmoid(e) + g = X @ U + h = f * g + i = h @ W + + ### Backpropagation chain rule + See our blog post for more details + + df = sigmoid(e) * (1 - f) + f + dC/dW = h.T @ dY + dC/dU = X.T @ (D @ W.T * f) + dC/dG = X.T @ (D @ W.T * df * g) + + ### Down projection LoRA weights + dC/dAw = dC/dW @ B.T + dC/dBw = A.T @ dC/dW + dC/dAw = h.T @ dY @ B.T + dC/dBw = A.T @ h.T @ dY + + ### Up projection LoRA weights + dC/dAu = X.T @ (D @ W.T * f) @ B.T + dC/dBu = A.T @ X.T @ (D @ W.T * f) + + ### Gate projection LoRA weights + dC/dAg = X.T @ (D @ W.T * df * g) @ B.T + dC/dBg = A.T @ X.T @ (D @ W.T * df * g) + + Don't forget to see our blog post for more details! + """ + @staticmethod + @torch_amp_custom_fwd + def forward(ctx, X: torch.Tensor, + gateupW, gateupW_quant, gateupA, gateupB, gateupS, + downW, downW_quant, downA, downB, downS, + _forward_function, _backward_function,): + dtype = X.dtype + + res_gateup_proj = matmul_lora( + X, gateupW, gateupW_quant, gateupA, gateupB, gateupS) + # e, g = torch.chunk(res_gateup_proj, 2, dim=-1) + e, g = torch.split( + res_gateup_proj, res_gateup_proj.size(-1)//2, dim=-1) + h = _forward_function(e, g) + i = matmul_lora(h, downW, downW_quant, downA, downB, downS) + + ctx.custom_saved_tensors = ( + gateupW, gateupW_quant, gateupS, + downW, downW_quant, downS, + _backward_function, + ) + ctx.save_for_backward(gateupA, gateupB, downA, downB, X, e, g) + return i + pass + + @staticmethod + @torch_amp_custom_bwd + def backward(ctx, dY: torch.Tensor): + gateupW, gateupW_quant, gateupS, downW, downW_quant, downS, \ + _backward_function = ctx.custom_saved_tensors + gateupA, gateupB, downA, downB, \ + X, e, g = ctx.saved_tensors + + gateupA, gateupB, downA, downB = \ + gateupA.t(), gateupB.t(), downA.t(), downB.t() + + batch, seq_len, hd = X.shape + dY = dY.view(-1, dY.shape[-1]) + X = X .view(-1, X .shape[-1]) + e = e .view(-1, e .shape[-1]) + g = g .view(-1, g .shape[-1]) + dtype = X.dtype + + DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) + DW, e, g = _backward_function(DW, e, g) + h, df, de = DW, e, g + + # Down projection LoRA weights + d_downA = h.t() @ (dY @ downB.t()) + d_downB = (downA.t() @ h.t()) @ dY + d_downA *= downS + d_downB *= downS + + # Gate_up projection LoRA weights + d_gateupA = X.t() @ (de @ gateupB.t()) + d_gateupB = (gateupA.t() @ X.t()) @ de + d_gateupA *= gateupS + d_gateupB *= gateupS + + # dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS) + # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS) + gateupW = fast_dequantize(gateupW.t(), gateupW_quant) + dX = de @ gateupW.t() + del gateupW + dX += de @ gateupB.to(dtype).t() @ (gateupS * gateupA.to(dtype).t()) + + # gateW, gateW_quant, gateA, gateB, gateS, + # upW, upW_quant, upA, upB, upS, + # downW, downW_quant, downA, downB, downS, + return dX.view(batch, seq_len, hd), \ + None, None, d_gateupA.t(), d_gateupB.t(), None, \ + None, None, d_downA.t(), d_downB.t(), None, \ + None, None, # _backward and _forward + pass + + +pass + + +def apply_lora_fuse_mlp_swiglu(self, X): + gateupW, gateupW_quant, gateupA, gateupB, gateupS = get_lora_parameters( + self.gate_up) + downW, downW_quant, downA, downB, downS = get_lora_parameters( + self.down_proj) + + out = LoRA_FUSEMLP.apply(X, + gateupW, gateupW_quant, gateupA, gateupB, gateupS, + downW, downW_quant, downA, downB, downS, + swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel,) + return out + + +pass diff --git a/ixformer_sdk/train/speedformer/layers/fast_lora/fast_lora_.py b/ixformer_sdk/train/speedformer/layers/fast_lora/fast_lora_.py new file mode 100644 index 00000000..95984a8f --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/fast_lora/fast_lora_.py @@ -0,0 +1,148 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# 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 torch +from .utils import ( + fast_dequantize, + QUANT_STATE, + get_lora_parameters, + matmul_lora, + torch_amp_custom_fwd, + torch_amp_custom_bwd, +) + + +class LoRA_MLP(torch.autograd.Function): + """ + ### LoRA weights + G = G + Ag @ Bg + U = U + Au @ Bu + W = W + Aw @ Bw + + ### SwiGLU(X) + e = X @ G + f = e * sigmoid(e) + g = X @ U + h = f * g + i = h @ W + + ### Backpropagation chain rule + See our blog post for more details + + df = sigmoid(e) * (1 - f) + f + dC/dW = h.T @ dY + dC/dU = X.T @ (D @ W.T * f) + dC/dG = X.T @ (D @ W.T * df * g) + + ### Down projection LoRA weights + dC/dAw = dC/dW @ B.T + dC/dBw = A.T @ dC/dW + dC/dAw = h.T @ dY @ B.T + dC/dBw = A.T @ h.T @ dY + + ### Up projection LoRA weights + dC/dAu = X.T @ (D @ W.T * f) @ B.T + dC/dBu = A.T @ X.T @ (D @ W.T * f) + + ### Gate projection LoRA weights + dC/dAg = X.T @ (D @ W.T * df * g) @ B.T + dC/dBg = A.T @ X.T @ (D @ W.T * df * g) + + Don't forget to see our blog post for more details! + """ + @staticmethod + @torch_amp_custom_fwd + def forward(ctx, X : torch.Tensor, + gateupW, gateupW_quant, gateupA, gateupB, gateupS, + downW, downW_quant, downA, downB, downS, + _forward_function, _backward_function,): + dtype = X.dtype + + res_gateup_proj = matmul_lora(X, gateupW, gateupW_quant, gateupA, gateupB, gateupS) + res_swiglu = _forward_function(res_gateup_proj) + res_mlp = matmul_lora(res_swiglu, downW, downW_quant, downA, downB, downS) + + ctx.custom_saved_tensors = ( + gateupW, gateupW_quant, gateupS, + downW, downW_quant, downS, + _backward_function, + ) + ctx.save_for_backward(gateupA, gateupB, downA, downB, X, res_gateup_proj, res_mlp) + return res_mlp + pass + + + @staticmethod + @torch_amp_custom_bwd + def backward(ctx, dY : torch.Tensor): + gateupW, gateupW_quant, gateupS, downW, downW_quant, downS, \ + _backward_function = ctx.custom_saved_tensors + gateupA, gateupB, downA, downB, \ + X, res_gateup_proj, res_mlp = ctx.saved_tensors + + gateupA, gateupB, downA, downB = \ + gateupA.t(), gateupB.t(), downA.t(), downB.t() + + batch, seq_len, hd = X.shape + dY = dY.view(-1, dY.shape[-1]) + X = X .view(-1, X .shape[-1]) + res_gateup_proj = res_gateup_proj.view(-1, res_gateup_proj.shape[-1]) + dtype = X.dtype + + D_swiglu = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) + DW, e, g = _backward_function(D_swiglu, res_gateup_proj) + h, df, de = DW, e, g + + # Down projection LoRA weights + d_downA = h.t() @ (dY @ downB.t()) + d_downB = (downA.t() @ h.t()) @ dY + d_downA *= downS + d_downB *= downS + + # Gate_up projection LoRA weights + d_gateupA = X.t() @ (de @ gateupB.t()) + d_gateupB = (gateupA.t() @ X.t()) @ de + d_gateupA *= gateupS + d_gateupB *= gateupS + + # dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS) + # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS) + + gateupW = fast_dequantize(gateupW.t(), gateupW_quant) + dX = de @ gateupW.t() + del gateupW + dX += de @ gateupB.to(dtype).t() @ (gateupS * gateupA.to(dtype).t()) + + # gateW, gateW_quant, gateA, gateB, gateS, + # upW, upW_quant, upA, upB, upS, + # downW, downW_quant, downA, downB, downS, + return dX.view(batch, seq_len, hd), \ + None, None, d_gateupA.t(), d_gateupB.t(), None, \ + None, None, d_downA.t(), d_downB.t(), None, \ + None, None, # _backward and _forward + pass +pass + + +from .swiglu_ import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel +def apply_lora_mlp_swiglu(self, X): + gateupW, gateupW_quant, gateupA, gateupB, gateupS = get_lora_parameters(self.gate_up) + downW, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj) + + out = LoRA_MLP.apply(X, + gateupW, gateupW_quant, gateupA, gateupB, gateupS, + downW, downW_quant, downA, downB, downS, + swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel,) + return out +pass \ No newline at end of file diff --git a/ixformer_sdk/train/speedformer/layers/fast_lora/swiglu.py b/ixformer_sdk/train/speedformer/layers/fast_lora/swiglu.py new file mode 100644 index 00000000..20791f6c --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/fast_lora/swiglu.py @@ -0,0 +1,106 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# 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 triton +import triton.language as tl +import torch + + +@triton.jit +def _fg_kernel(e, g, h, n_elements, BLOCK_SIZE: tl.constexpr,): + block_idx = tl.program_id(0) + offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + e_row = tl.load(e + offsets, mask=mask, other=0).to(tl.float32) + g_row = tl.load(g + offsets, mask=mask, other=0) # .to(tl.float32) + + # f = e * sigmoid(e) + f_row = e_row * tl.sigmoid(e_row) # e_row / (1 + tl.exp(-e_row)) + f_row = f_row.to(g_row.dtype) # Exact copy from HF + # h = f * g + h_row = f_row * g_row + + # Store h + tl.store(h + offsets, h_row, mask=mask) + + +pass + + +def swiglu_fg_kernel(e, g): + batch, seq_len, hd = e.shape + n_elements = e.numel() + h = torch.empty((batch, seq_len, hd), dtype=e.dtype, device="cuda:0") + def grid(meta): return (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + _fg_kernel[grid](e, g, h, n_elements, BLOCK_SIZE=1024,) + return h + + +pass + + +@triton.jit +def _DWf_DW_dfg_kernel(DW, e, g, n_elements, BLOCK_SIZE: tl.constexpr,): + """ + e = e.float() + se = 1.0 / (1.0 + torch.exp(-e)) + f = (se * e).to(dtype) + h = f * g + df = DW * f + dg = DW * g + de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype) + """ + block_idx = tl.program_id(0) + offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + DW_row = tl.load(DW + offsets, mask=mask, other=0) # .to(tl.float32) + e_row = tl.load(e + offsets, mask=mask, other=0).to(tl.float32) + g_row = tl.load(g + offsets, mask=mask, other=0) # .to(tl.float32) + + # e = e.float() + # se = 1.0 / (1.0 + torch.exp(-e)) + se_row = tl.sigmoid(e_row) # 1.0 / (1.0 + tl.exp(-e_row)) + # f = (se * e).to(dtype) + f_row = se_row * e_row + f_row = f_row.to(DW_row.dtype) + # h = f * g + h_row = f_row * g_row + # df = DW * f + df_row = DW_row * f_row + # dg = DW * g + dg_row = DW_row * g_row + # de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype) + de_row = dg_row.to(tl.float32) * se_row * (1.0 + e_row * (1.0 - se_row)) + de_row = de_row.to(DW_row.dtype) + + # Store derivatives in buffers + tl.store(DW + offsets, h_row, mask=mask) # h = f * g + tl.store(e + offsets, df_row, mask=mask) # df = DW * f + tl.store(g + offsets, de_row, mask=mask) # de + + +pass + + +def swiglu_DWf_DW_dfg_kernel(DW, e, g): + batch_seq_len, hd = e.shape + n_elements = e.numel() + def grid(meta): return (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + _DWf_DW_dfg_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE=1024,) + return DW, e, g + + +pass diff --git a/ixformer_sdk/train/speedformer/layers/fast_lora/swiglu_.py b/ixformer_sdk/train/speedformer/layers/fast_lora/swiglu_.py new file mode 100644 index 00000000..140544cc --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/fast_lora/swiglu_.py @@ -0,0 +1,102 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# 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 triton +import triton.language as tl +import torch + + +@triton.jit +def _fg_kernel(x, h, hd, BLOCK_SIZE : tl.constexpr,): + block_idx = tl.program_id(0) + offsets0 = block_idx*2*hd + tl.arange(0, BLOCK_SIZE) + offsets1 = block_idx*2*hd + hd + tl.arange(0, BLOCK_SIZE) + mask = offsets0 < hd + + e_row = tl.load(x + offsets0, mask = mask, other = 0).to(tl.float32) + g_row = tl.load(x + offsets1, mask = mask, other = 0)#.to(tl.float32) + + # f = e * sigmoid(e) + f_row = e_row * tl.sigmoid(e_row) # e_row / (1 + tl.exp(-e_row)) + f_row = f_row.to(g_row.dtype) # Exact copy from HF + # h = f * g + h_row = f_row * g_row + + # Store h + tl.store(h + offsets0, h_row, mask = mask) +pass + + +def swiglu_fg_kernel(x): + batch, seq_len, hdx2 = x.shape + hd = hdx2 // 2 + n_rows = batch * seq_len + BLOCK_SIZE = triton.next_power_of_2(hd) + h = torch.empty((batch, seq_len, hd), dtype = x.dtype, device = "cuda:0") + + _fg_kernel[n_rows,](x, h, hd, BLOCK_SIZE=BLOCK_SIZE) + return h +pass + + +@triton.jit +def _DWf_DW_dfg_kernel(DW, x, hd, BLOCK_SIZE : tl.constexpr,): + """ + e = e.float() + se = 1.0 / (1.0 + torch.exp(-e)) + f = (se * e).to(dtype) + h = f * g + df = DW * f + dg = DW * g + de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype) + """ + block_idx = tl.program_id(0) + offsets0 = block_idx*hd*2 + tl.arange(0, BLOCK_SIZE) + offsets1 = block_idx*hd*2 + hd + tl.arange(0, BLOCK_SIZE) + mask = BLOCK_SIZE < hd + + DW_row = tl.load(DW + offsets0, mask = mask, other = 0)#.to(tl.float32) + e_row = tl.load(x + offsets0, mask = mask, other = 0).to(tl.float32) + g_row = tl.load(x + offsets1, mask = mask, other = 0)#.to(tl.float32) + + # e = e.float() + # se = 1.0 / (1.0 + torch.exp(-e)) + se_row = tl.sigmoid(e_row) # 1.0 / (1.0 + tl.exp(-e_row)) + # f = (se * e).to(dtype) + f_row = se_row * e_row + f_row = f_row.to(DW_row.dtype) + # h = f * g + h_row = f_row * g_row + # df = DW * f + df_row = DW_row * f_row + # dg = DW * g + dg_row = DW_row * g_row + # de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype) + de_row = dg_row.to(tl.float32) * se_row * (1.0 + e_row * (1.0 - se_row)) + de_row = de_row.to(DW_row.dtype) + + # Store derivatives in buffers + tl.store(DW + offsets0, h_row, mask = mask) # h = f * g + tl.store(x + offsets0, df_row, mask = mask) # df = DW * f + tl.store(x + offsets1, de_row, mask = mask) # de +pass + + +def swiglu_DWf_DW_dfg_kernel(DW, x): + batch_seq_len, hdx2 = x.shape + hd = hdx2 // 2 + BLOCK_SIZE = triton.next_power_of_2(hd) + _DWf_DW_dfg_kernel[batch_seq_len, ](DW, x, hd, BLOCK_SIZE=BLOCK_SIZE,) + return DW, x +pass diff --git a/ixformer_sdk/train/speedformer/layers/fast_lora/utils.py b/ixformer_sdk/train/speedformer/layers/fast_lora/utils.py new file mode 100644 index 00000000..24cb3bef --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/fast_lora/utils.py @@ -0,0 +1,195 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# 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 ctypes +import bitsandbytes as bnb +from packaging.version import Version +import torch +import triton +MAX_FUSED_SIZE = 65536 +next_power_of_2 = triton.next_power_of_2 + +# torch.cuda.amp.custom_fwd is deprecated >= 2.4 +if Version(torch.__version__) < Version("2.4.0"): + torch_amp_custom_fwd = torch.cuda.amp.custom_fwd + torch_amp_custom_bwd = torch.cuda.amp.custom_bwd +else: + torch_amp_custom_fwd = torch.amp.custom_fwd(device_type="cuda") + torch_amp_custom_bwd = torch.amp.custom_bwd(device_type="cuda") +pass + + +# tl.math.tanh now is libdevice.tanh +if Version(triton.__version__) >= Version("3.0.0"): + from triton.language.extra import libdevice + triton_tanh = libdevice.tanh +else: + import triton.language as tl + triton_tanh = tl.math.tanh +pass + + +def calculate_settings(n): + BLOCK_SIZE = next_power_of_2(n) + if BLOCK_SIZE > MAX_FUSED_SIZE: + raise RuntimeError(f"Cannot launch Triton kernel since n = {n} exceeds " + f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.") + num_warps = 4 + if BLOCK_SIZE >= 32768: + num_warps = 32 + elif BLOCK_SIZE >= 8192: + num_warps = 16 + elif BLOCK_SIZE >= 2048: + num_warps = 8 + return BLOCK_SIZE, num_warps + + +pass + + +get_ptr = bnb.functional.get_ptr +cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32 +cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4 +cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4 + + +def QUANT_STATE(W): + return getattr(W, "quant_state", None) + + +pass + + +def get_lora_parameters(proj): + # For DPO or disabled adapters + base_layer = (proj.base_layer if hasattr(proj, "base_layer") else proj) + W = base_layer.weight + + if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged: + return W, QUANT_STATE(W), None, None, None + pass + + active_adapter = proj.active_adapters[0] if \ + hasattr(proj, "active_adapters") else proj.active_adapter + A = proj.lora_A[active_adapter].weight + B = proj.lora_B[active_adapter].weight + s = proj.scaling[active_adapter] + return W, QUANT_STATE(W), A, B, s + + +pass + + +def get_lora_parameters_bias(proj): + # For DPO or disabled adapters + base_layer = (proj.base_layer if hasattr(proj, "base_layer") else proj) + W = base_layer.weight + bias = base_layer.bias + + if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged: + return W, QUANT_STATE(W), None, None, None, bias + pass + + active_adapter = proj.active_adapters[0] if \ + hasattr(proj, "active_adapters") else proj.active_adapter + A = proj.lora_A[active_adapter].weight + B = proj.lora_B[active_adapter].weight + s = proj.scaling[active_adapter] + return W, QUANT_STATE(W), A, B, s, bias + + +pass + + +def fast_dequantize(W, quant_state=None, out=None): + if quant_state is None: + return W + if type(quant_state) is not list: + # New quant_state as a class + # https://github.com/TimDettmers/bitsandbytes/pull/763/files + absmax = quant_state.absmax + shape = quant_state.shape + dtype = quant_state.dtype + blocksize = quant_state.blocksize + offset = quant_state.offset + state2 = quant_state.state2 + absmax2 = state2.absmax + code2 = state2.code + blocksize2 = state2.blocksize + else: + # Old quant_state as a list of lists + absmax, shape, dtype, blocksize, compressed_stats, _, _ = quant_state + offset, state2 = compressed_stats + absmax2, code2, blocksize2, _, _, _, _ = state2 + pass + + # Create weight matrix + if out is None: + out = torch.empty(shape, dtype=dtype, device="cuda:0") + else: + assert (out.shape == shape) + assert (out.dtype == dtype) + + # NF4 dequantization of statistics + n_elements_absmax = absmax.numel() + out_absmax = torch.empty( + n_elements_absmax, dtype=torch.float32, device="cuda:0") + + # Do dequantization + ptr_out_absmax = get_ptr(out_absmax) + cdequantize_blockwise_fp32( + get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax, + ctypes.c_int(blocksize2), ctypes.c_int(n_elements_absmax) + ) + out_absmax += offset + + fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \ + cdequantize_blockwise_bf16_nf4 + fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out), + ctypes.c_int(blocksize), ctypes.c_int(out.numel())) + + # Careful returning transposed data + is_transposed = (True if W.shape[0] == 1 else False) + return out.t() if is_transposed else out + + +pass + + +def matmul_lora(X, W, W_quant, A, B, s, out=None): + dtype = X.dtype + W = fast_dequantize(W.t(), W_quant) + + if X.dim() == 3: + batch, seq_len, d = X.shape + X = X.view(-1, X.shape[-1]) + reshape = True + else: + reshape = False + pass + + out = torch.matmul(X, W, out=out) + if W_quant is not None: + del W + + if A is not None: + # LoRA is enabled + A, B = A.t(), B.t() + out += (X @ A.to(dtype)) @ (s * B.to(dtype)) + pass + + return out.view(batch, seq_len, -1) if reshape else out + + +pass diff --git a/ixformer_sdk/train/speedformer/layers/gpt2/__init__.py b/ixformer_sdk/train/speedformer/layers/gpt2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/layers/gpt2/attention.py b/ixformer_sdk/train/speedformer/layers/gpt2/attention.py new file mode 100644 index 00000000..4ae1e406 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/gpt2/attention.py @@ -0,0 +1,45 @@ +import torch +import os +from einops import rearrange +from flash_attn import flash_attn_varlen_func + + +@staticmethod +def replace_flash_attn_forward(self, q, k, v, attention_mask, query_length, dropout=0.0, softmax_scale=None): + + # flash-attn(ixdnn)存在gpt2(118M,338M,738M) shape没适配，只能采用普通版本 + assert os.getenv('ENABLE_FLASH_ATTENTION_WITH_IXDNN', "1") == '0', "flash-attn should not be use ixdnn version, please set variables" \ + " in shell \"export ENABLE_FLASH_ATTENTION_WITH_IXDNN=0 \" " + assert all((i.dtype in [torch.float16, torch.bfloat16] for i in (q, k, v))) + assert all((i.is_cuda for i in (q, k, v))) + + batch_size, seqlen_q = q.shape[0], q.shape[1] + seqlen_k = k.shape[1] + + q, k, v = [rearrange(x, 'b s ... -> (b s) ...') for x in [q, k, v]] + cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, + device=q.device) + + if query_length != 1: + # during training q,k,v always have same seqlen + assert seqlen_k == seqlen_q + + is_causal = self.is_causal + cu_seqlens_k = cu_seqlens_q + dropout_p = dropout + else: + # turn off FA causal mask after first inference autoregressive iteration + # only on first autoregressive step q,k,v have same seqlen + is_causal = seqlen_q == seqlen_k + cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, + device=q.device) + dropout_p = 0 + + output = flash_attn_varlen_func( + q, k, v, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k, + dropout_p, + softmax_scale=softmax_scale, causal=is_causal + ) + # print(f"{output}") + output = rearrange(output, '(b s) ... -> b s ...', b=batch_size) + return output diff --git a/ixformer_sdk/train/speedformer/layers/lazy/__init__.py b/ixformer_sdk/train/speedformer/layers/lazy/__init__.py new file mode 100644 index 00000000..c6b813c5 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/lazy/__init__.py @@ -0,0 +1,6 @@ +from .lazy_init import LazyInitContext, LazyTensor + +__all__ = [ + "LazyInitContext", + "LazyTensor", +] diff --git a/ixformer_sdk/train/speedformer/layers/lazy/construction.py b/ixformer_sdk/train/speedformer/layers/lazy/construction.py new file mode 100644 index 00000000..6764eaf7 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/lazy/construction.py @@ -0,0 +1,87 @@ +from contextlib import contextmanager +from typing import Callable, Dict, Tuple + +import torch + +__all__ = [ + "_LEGACY_TENSOR_CONSTRUCTOR", + "_NO_META_FACTORY", + "_NORMAL_FACTORY", + "ConstructorManager", +] + +# reference: https://pytorch.org/cppdocs/notes/tensor_creation.html +_NORMAL_FACTORY = [ + "arange", + "full", + "empty", + "linspace", + "logspace", + "ones", + "rand", + "randn", + "randint", + "randperm", + "zeros", + "tensor", +] + +# factory function that does not support meta tensor backend +_NO_META_FACTORY = [ + "eye", +] + +_LEGACY_TENSOR_CONSTRUCTOR = { + "FloatTensor": torch.float, + "DoubleTensor": torch.double, + "HalfTensor": torch.half, + "BFloat16Tensor": torch.bfloat16, + "ByteTensor": torch.uint8, + "CharTensor": torch.int8, + "ShortTensor": torch.short, + "IntTensor": torch.int, + "LongTensor": torch.long, + "BoolTensor": torch.bool, +} + + +class ConstructorManager: + # function name: (new, old) + overwrites: Dict[str, Tuple[Callable, Callable]] = {} + changed: bool = False + + @staticmethod + def apply(overwrites: Dict[Callable, Callable]): + ConstructorManager.overwrites.clear() + ConstructorManager.overwrites.update(overwrites) + ConstructorManager.redo() + + @staticmethod + def undo(): + assert ConstructorManager.changed, "No constructor change to undo" + for name, (new, old) in ConstructorManager.overwrites.items(): + setattr(torch, name, old) + ConstructorManager.changed = False + + @staticmethod + def redo(): + assert not ConstructorManager.changed, "Constructor already changed" + for name, (new, old) in ConstructorManager.overwrites.items(): + setattr(torch, name, new) + ConstructorManager.changed = True + + @staticmethod + @contextmanager + def disable(): + enabled = ConstructorManager.changed + if enabled: + ConstructorManager.undo() + yield + if enabled: + ConstructorManager.redo() + + @staticmethod + def clear(): + if ConstructorManager.changed: + ConstructorManager.undo() + ConstructorManager.overwrites.clear() diff --git a/ixformer_sdk/train/speedformer/layers/lazy/lazy_init.py b/ixformer_sdk/train/speedformer/layers/lazy/lazy_init.py new file mode 100644 index 00000000..064f3968 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/lazy/lazy_init.py @@ -0,0 +1,669 @@ +from types import MethodType +from typing import Callable, Optional, Union + +import torch +import torch.nn as nn +from packaging import version +from torch import Tensor +from torch.nn import Parameter +from torch.utils._pytree import tree_map + +from ixformer.train.speedformer.layers.lazy.construction import ConstructorManager +from ixformer.train.speedformer.layers.lazy.pretrained import PretrainedManager + +# reference: https://pytorch.org/cppdocs/notes/tensor_creation.html +_NORMAL_FACTORY = [ + "arange", + "full", + "empty", + "linspace", + "logspace", + "ones", + "rand", + "randn", + "randint", + "randperm", + "zeros", + "tensor", +] + +# factory function that does not support meta tensor backend +_NO_META_FACTORY = [ + "eye", +] + +_EARLY_MATERIALIZED_OPS = ["__getitem__", "split"] + +# If your intent is to change the metadata of a Tensor (such as sizes / strides / storage / storage_offset) +# without autograd tracking the change, remove the .data / .detach() call and wrap the change in a `with torch.no_grad():` block. +# These ops cannot be unwrapped using .data +_CHANGE_META_OPS = ["_cudnn_rnn_flatten_weight", + "requires_grad_", "__get__", "__set__", "numel", "size", "dim"] + +# These ops is not related to tensor value and should not be rerun +_NO_RERUN_OPS = ["__get__", "numel", "size", "dim"] + +_LEGACY_TENSOR_CONSTRUCTOR = { + "FloatTensor": torch.float, + "DoubleTensor": torch.double, + "HalfTensor": torch.half, + "BFloat16Tensor": torch.bfloat16, + "ByteTensor": torch.uint8, + "CharTensor": torch.int8, + "ShortTensor": torch.short, + "IntTensor": torch.int, + "LongTensor": torch.long, + "BoolTensor": torch.bool, +} + +# These ops have at least one lazy tensor argument and maybe a scalar argument +# scalar value should be converted to meta tensor +# this is a hack for torch 2.0 +_EXPAND_SCALAR_OPS = [ + "where", + "clamp", + "clamp_min", + "clamp_max", + "clamp_", + "clamp_min_", + "clamp_max_", +] +_old_tensor_factory = torch.tensor + +_EMPTY_DATA = torch.empty(0) + + +class _MyTensor(Tensor): + """This class is only for correctness verification.""" + + _pre_op_fn: Callable[["LazyTensor"], None] = lambda *args: None + + default_device: Optional[torch.device] = None + + def __new__(cls, func, *args, concrete_data=None, **kwargs) -> "_MyTensor": + cls._pre_op_fn() + if concrete_data is not None: + # uniform api as LazyTensor + data = concrete_data + else: + kwargs["device"] = cls.default_device + data = func(*args, **kwargs) + return Tensor._make_subclass(cls, data, require_grad=data.requires_grad) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + cls._pre_op_fn() + return super().__torch_function__(func, types, args, kwargs) + + +def _data_tolist(tensor: torch.Tensor) -> list: + """tolist() method is not allowed for a subclass of tensor. Tensor.data returns a Tensor.""" + return tensor.data.tolist() + + +def _convert_cls(tensor: "LazyTensor", target: torch.Tensor) -> torch.Tensor: + """Convert a lazy tensor's class to target's class, with target's data. + + The reason why we change the class of a lazy tensor in-place is that this can easily handle shared modules/parameters, which is common in huggingface models. + If we create a new tensor and update the module by ``setattr(module, name, param)``, the shared parameters will not be updated. And we have to track all shared parameters and update them manually. + + Args: + tensor (LazyTensor): the LazyTensor to be converted + target (torch.Tensor): target tensor + + Returns: + torch.Tensor: the converted tensor + """ + cls_to_become = Parameter if isinstance( + tensor, Parameter) else torch.Tensor + tensor.__class__ = cls_to_become + if cls_to_become is Parameter: + # to fit UninitializedParameter + delattr(tensor, "_is_param") + tensor.data = target + tensor.requires_grad = target.requires_grad + # subclass of torch.Tensor does not have tolist() method + # overwrite this method after materialization or distribution + tensor.tolist = MethodType(_data_tolist, tensor) + return tensor + + +class LazyTensor(torch.Tensor): + """A naive implementation of LazyTensor (https://arxiv.org/pdf/2102.13267.pdf). + + Usage: + 1. Use ``LazyTensor`` instead of ``torch.Tensor``. + >>> x = LazyTensor(torch.zeros, 2, 3) + >>> x += 1 + >>> y = x * x + >>> y = y.cuda().half() + >>> y[0, 0] = 0 + >>> y = y.materialize() # materialize the tensor + >>> print(y) + tensor([[0., 1., 1.], + [1., 1., 1.]], device='cuda:0', dtype=torch.float16) + + Warnings: + 1. Cases that ``LazyTensor`` can't deal with. + >>> x = LazyTensor(torch.ones, 2, 3) + >>> x[0, 0] = -x[0, 0] # this will cause infinite recursion + >>> y = x.clone() + >>> x.add_(1) # modifying origin tensor after cloning leads to wrong materialization + >>> z = x.tolist() + >>> x.zeros_() # modifying origin tensor after cloning tolist is not allowed + >>> nn.utils.weight_norm(self.conv, name="weight", dim=2) # applying weight norm on a lazy tensor is not allowed + + + 2. Cases that ``LazyTensor`` becomes eager (early materialization). + >>> b = a[:, 2:] # get a slice of a lazy tensor triggers early materialization + >>> chunks = a.split(3) # this also triggers early materialization + >>> x.data = torch.rand(2, 3) # directly setting data of a lazy tensor triggers early materialization + + """ + + _repr = True + _meta_data: Optional[torch.Tensor] = None # shape, dtype, device + _pre_op_fn: Callable[["LazyTensor"], None] = lambda *args: None + + default_device: Optional[torch.device] = None + _device: torch.device # fake device of mate tensor + + @staticmethod + def __new__(cls, func, *args, meta_data=None, concrete_data=None, **kwargs): + # tips for torch 2.0: + # torch 2.0 disables torch dispatch for subclass of tensor + # MetaTensor is cannot be used + # Now lazy tensor contains device injection and meta tensor + if concrete_data is not None: + # some ops don't support meta backend and should have concrete data + elem = concrete_data + else: + if meta_data is None: + with ConstructorManager.disable(): + # to disable create lazy tensor in inner ops, this is a hack for torch 2.0 + meta_data = func(*args, **{**kwargs, "device": "meta"}) + elem = meta_data + # As a meta tensor cannot be modified __class__ to torch.Tensor, we should use an empty real tensor here + r = torch.Tensor._make_subclass( + cls, _EMPTY_DATA, require_grad=elem.requires_grad) + r._meta_data = meta_data + + return r + + def __init__(self, func, *args, meta_data=None, concrete_data=None, **kwargs): + self._device = torch.device(kwargs.get("device", None) or "cpu") + if func.__name__ in _NORMAL_FACTORY: + kwargs = {**kwargs, "device": LazyTensor.default_device} + self._factory_method = (func, args, kwargs) # (func, args, kwargs) + self._op_buffer = [] # (func, args, kwargs, replace) + # materialized data + self._materialized_data: Optional[torch.Tensor] = concrete_data + + @property + def device(self) -> torch.device: + return self._materialized_data.device if self._materialized_data is not None else self._device + + def __repr__(self): + return f"LazyTensor(..., size={tuple(self.shape)}, device='{self.device}', dtype={self.dtype})" + + def materialize(self) -> torch.Tensor: + """Materialize the ``LazyTensor`` to ``torch.Tensor`` by modifying __class__ (inplace). + + Returns: + torch.Tensor: The materialized tensor (self). + """ + target = self._materialize_data() + self.clean() + return _convert_cls(self, target) + + def clean(self) -> None: + """Clean all stored operations, meta data and materialized data, which prevents memory leaking. This should be called after all tensors are materialized.""" + delattr(self, "_factory_method") + delattr(self, "_op_buffer") + delattr(self, "_materialized_data") + delattr(self, "_meta_data") + + @staticmethod + def _replace_with_materialized(x): + if isinstance(x, LazyTensor): + return x._materialize_data() + return x + + def _materialize_data(self) -> torch.Tensor: + # self._materialized_data should be generated after the first call of this function + if self._materialized_data is None: + # apply factory method + func, args, kwargs = self._factory_method + # apply cached sequence + self._pre_op_fn() + + init_val = func( + *tree_map(self._replace_with_materialized, args), **tree_map(self._replace_with_materialized, kwargs) + ) + + self._materialized_data = self._rerun_ops(init_val) + return self._materialized_data + + def _rerun_ops(self, target=None) -> torch.Tensor: + """Do lazy execution by rerunning all (stored) related operations. + + Args: + target (torc.Tensor, optional): Intial value of the target tensor (self). Defaults to None. + """ + + def replace(x): + if x is self: + return target + elif isinstance(x, LazyTensor): + return x._materialize_data() + return x + + packed = None + + for func, args, kwargs in self._op_buffer: + if func == torch.Tensor.requires_grad_: + packed = func, args, kwargs # requires grad should be set at last + else: + self._pre_op_fn() + o = func(*tree_map(replace, args), **tree_map(replace, kwargs)) + # if func returns non-Tensor, discard the value + target = o if isinstance(o, torch.Tensor) else target + + # super-dainiu: set requires_grad after all inplace-ops are done + if packed is not None: + func, args, kwargs = packed + func(*tree_map(replace, args), **tree_map(replace, kwargs)) + + return target + + # cache everything with __torch_function__ + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + if func.__name__ in _EARLY_MATERIALIZED_OPS: + # These OPs cannot be lazy and related tensors should be early materialized + tree_map(cls._replace_with_materialized, args) + tree_map(cls._replace_with_materialized, kwargs) + is_inplace: bool = ( + func.__name__.endswith("_") + and not (func.__name__.endswith("__")) + or func.__name__ in ("__setitem__", "__set__") + ) + + is_change_meta_op: bool = func.__name__ in _CHANGE_META_OPS + + if isinstance(func, torch._C.ScriptMethod): + # FIXME(ver217): torch script functions are not verified + + target = None + + def unwrap(x): + if isinstance(x, LazyTensor): + return x._meta_data + return x + + target: LazyTensor = args[0].clone() + target._op_buffer.append((func, args, kwargs)) + target._meta_data = getattr(target._meta_data, func.name)( + *tree_map(unwrap, args[1:]), **tree_map(unwrap, kwargs) + ) + return target + else: + meta_to_lazy = {} + + def unwrap(x): + if isinstance(x, LazyTensor): + if x._materialized_data is not None: + # for early materialized tensor, use its materialized data directly + return x._materialized_data if is_change_meta_op else x._materialized_data.data + t = x if is_inplace else x.clone() + if func.__name__ not in _NO_RERUN_OPS: + t._op_buffer.append((func, args, kwargs)) + meta = x._meta_data if is_change_meta_op else x._meta_data.data + meta_to_lazy[meta] = t + return meta + elif ( + version.parse(torch.__version__) >= version.parse("2.0.0") + and func.__name__ in _EXPAND_SCALAR_OPS + and not isinstance(x, torch.Tensor) + ): + return _old_tensor_factory(x, device="meta") + return x + + def wrap(y, i=None): + if isinstance(y, torch.Tensor): + if y.is_meta: + if y in meta_to_lazy: + # inplace op, just return origin lazy tensor + return meta_to_lazy[y] + else: + # out of place op, create new lazy tensor + fn = lambda *a, **kw: func(*a, ** + kw) if i is None else func(*a, **kw)[i] + fn.__name__ = func.__name__ + lazy_y = LazyTensor( + fn, *args, meta_data=y, **kwargs) + return lazy_y + else: + # for early materialized tensor + return LazyTensor(lambda: None, concrete_data=y) + return y + + cls._pre_op_fn() + with ConstructorManager.disable(): + # to disable create lazy tensor in inner ops, this is a hack for torch 2.0 + o = func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs)) + if isinstance(o, (tuple, list)): + return type(o)(wrap(y, i=i) for i, y in enumerate(o)) + return wrap(o) + + def to(self, *args, **kwargs) -> torch.Tensor: + if self._materialized_data is not None: + return LazyTensor(lambda: None, concrete_data=self._materialized_data.to(*args, **kwargs)) + + device = None + + def replace(x): + nonlocal device + if isinstance(x, (str, int, torch.device)) and not isinstance(x, bool): + device = x + return torch.device("meta") + return x + + meta_data = self._meta_data.to( + *tree_map(replace, args), **tree_map(replace, kwargs)) + + if meta_data is self._meta_data and device == self.device: + return self + + def factory_fn(t: torch.Tensor, **kw): + return t.to(*args, **kwargs) + + return LazyTensor(factory_fn, self, meta_data=meta_data, device=device) + + def cpu(self, memory_format: torch.memory_format = torch.preserve_format): + return self.to(device=torch.device("cpu"), memory_format=memory_format) + + def cuda(self, device=None, non_blocking=False, memory_format: torch.memory_format = torch.preserve_format): + device = torch.device(device or "cuda") + return self.to(device=device, non_blocking=non_blocking, memory_format=memory_format) + + def clone(self) -> "LazyTensor": + def factory_fn(t: torch.Tensor, **kw): + # if self is materialized, return self + return t.clone() + + target = LazyTensor(factory_fn, self, meta_data=self._meta_data) + + return target + + def detach(self) -> Tensor: + return self + + def __deepcopy__(self, memo): + if not self.is_leaf: + raise RuntimeError( + "Only Tensors created explicitly by the user " + "(graph leaves) support the deepcopy protocol at the moment" + ) + if id(self) in memo: + return memo[id(self)] + + def factory_fn(t: torch.Tensor, **kw): + # if self is materialized, return self + return _copy_tensor(t, t.requires_grad) + + if self._materialized_data is not None: + # self is early materialized + copied = _copy_tensor(self._materialized_data, self.requires_grad) + target = LazyTensor(lambda: None, concrete_data=copied) + else: + target = LazyTensor(factory_fn, self, meta_data=self._meta_data) + + if isinstance(self, Parameter): + # hack isinstance check of parameter + target._is_param = True + + memo[id(self)] = target + return target + + @property + def data(self): + return self + + @data.setter + def data(self, other: "LazyTensor"): + """This is sightly different from oringinal `data` setter. + + E.g.: + >>> a = torch.randn(3, 3) # a is a Tensor + >>> b = torch.rand(2, 2) + >>> a.data = b + >>> b.add_(1) # this will affect a + >>> x = torch.randn(3, 3) # x is a LazyTensor + >>> y = torch.rand(2, 2) # y is a LazyTensor + >>> x.data = y + >>> y.add_(1) # this will not affect x + + """ + if other is self: + return + + def replace(x): + if x is other: + return self + return x + + for func, args, kwargs in [other._factory_method, *other._op_buffer]: + self._op_buffer.append( + (func, tree_map(replace, args), tree_map(replace, kwargs))) + + def tolist(self) -> list: + # Though self.__class__ is modified to torch.Tensor, in C++ side, it is still a subclass of torch.Tensor + # And subclass of torch.Tensor does not have tolist() method + t = self._materialize_data() + return t.tolist() + + def __hash__(self): + return id(self) + + def __rpow__(self, other): + dtype = torch.result_type(self, other) + return torch.tensor(other, dtype=dtype, device=self.device) ** self + + +class LazyInitContext: + """Context manager for lazy initialization. Enables initializing the model without allocating real memory. + + Args: + tensor_cls (Union[_MyTensor, LazyTensor], optional): This is only for test. Defaults to LazyTensor. + default_device (Optional[Union[torch.device, str, int]], optional): Defalt device for initialization. + If it's cuda, initilization will be accelerated, but cuda memory will be allocated. By default, it's cpu. + Defaults to None. + """ + + _replaced: bool = False + + def __init__( + self, + tensor_cls: Union[_MyTensor, LazyTensor] = LazyTensor, + default_device: Optional[Union[torch.device, str, int]] = None, + ): + assert tensor_cls is LazyTensor or tensor_cls is _MyTensor + self.tensor_cls = tensor_cls + self.old_default_device = LazyTensor.default_device + self.default_device = default_device + + def __enter__(self): + if LazyInitContext._replaced: + raise RuntimeError(f"LazyInitContext is not reentrant") + LazyInitContext._replaced = True + self.old_default_device = self.tensor_cls.default_device + self.tensor_cls.default_device = self.default_device + + def wrap_factory_method(target): + # factory functions (eg. torch.empty()) + def wrapper(*args, **kwargs): + return self.tensor_cls(target, *args, **kwargs) + + return wrapper, target + + def wrap_factory_like_method(orig_target, target): + # factory_like functions (eg. torch.empty_like()) + def wrapper(*args, **kwargs): + orig_t = args[0] + return self.tensor_cls( + orig_target, *orig_t.shape, *args[1:], device=orig_t.device, dtype=orig_t.dtype, **kwargs + ) + + return wrapper, target + + def wrap_legacy_constructor(target, dtype): + # legacy constructor (e.g. torch.LongTensor()) + def wrapper(*args, **kwargs): + if len(args) == 1 and isinstance(args[0], torch.Tensor): + # (Tensor other) + return args[0] + elif len(args) == 1: + # (object data, *, torch.device device) + kwargs = {**kwargs, "dtype": dtype} + replaced, orig = self.overrides["tensor"] + return replaced(*args, **kwargs) + elif _is_int_tuple(args): + # (tuple of ints size, *, torch.device device) + kwargs = {**kwargs, "dtype": dtype} + replaced, orig = self.overrides["empty"] + return replaced(*args, **kwargs) + else: + raise TypeError( + f"new() received an invalid combination of arguments - got {tuple(type(x) for x in args)}, but expected one of:\n * (Tensor other)\n * (tuple of ints size, *, torch.device device)\n * (object data, *, torch.device device)" + ) + + return wrapper, target + + def wrap_no_meta_factory(target): + # factory functions which don't support meta tensor backend + def wrapper(*args, **kwargs): + tensor = target(*args, **kwargs) + return self.tensor_cls(lambda: None, concrete_data=tensor) + + return wrapper, target + + overrides = { + target: wrap_factory_method(getattr(torch, target)) + for target in _NORMAL_FACTORY + if callable(getattr(torch, target, None)) + } + + overrides.update( + { + target + "_like": wrap_factory_like_method(getattr(torch, target), getattr(torch, target + "_like")) + for target in _NORMAL_FACTORY + if callable(getattr(torch, target + "_like", None)) + } + ) + + overrides.update( + { + target: wrap_legacy_constructor(getattr(torch, target), dtype) + for target, dtype in _LEGACY_TENSOR_CONSTRUCTOR.items() + if callable(getattr(torch, target, None)) + } + ) + + overrides.update( + { + target: wrap_no_meta_factory(getattr(torch, target)) + for target in _NO_META_FACTORY + if callable(getattr(torch, target, None)) + } + ) + + ConstructorManager.apply(overrides) + PretrainedManager.inject() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.tensor_cls.default_device = self.old_default_device + LazyInitContext._replaced = False + ConstructorManager.clear() + PretrainedManager.recover() + + @staticmethod + def materialize(module: nn.Module, verbose: bool = False) -> nn.Module: + """Initialize all ``Parameter`` from ``LazyTensor``. This function will modify the module in-place. + + Args: + module (nn.Module): Target ``nn.Module`` + verbose (bool): Whether to print lazy initialization rate. Defaults to False. + """ + + def apply_fn(name: str, p: LazyTensor): + p.materialize() + + return _apply_to_lazy_module(module, apply_fn, verbose) + + +def _apply_to_lazy_module( + module: nn.Module, apply_fn: Callable[[str, torch.Tensor], None], verbose: bool = False +) -> nn.Module: + if verbose: + # verbose info + param_cnt = 0 + param_lazy_cnt = 0 + buf_cnt = 0 + buf_lazy_cnt = 0 + total_numel = 0 + non_lazy_numel = 0 + + for name, p in module.named_parameters(): + if verbose: + param_cnt += 1 + total_numel += p.numel() + if getattr(p, "_materialized_data", False) is None: + # if no _materialized_data attr, the tensor is not lazy + param_lazy_cnt += 1 + else: + non_lazy_numel += p.numel() + if isinstance(p, LazyTensor): + apply_fn(name, p) + + for name, buf in module.named_buffers(): + if verbose: + buf_cnt += 1 + total_numel += buf.numel() + if getattr(buf, "_materialized_data", False) is None: + # if no _materialized_data attr, the tensor is not lazy + buf_lazy_cnt += 1 + else: + non_lazy_numel += buf.numel() + if isinstance(buf, LazyTensor): + apply_fn(name, buf) + + # if verbose: + # non_lazy_numel_ratio = non_lazy_numel / total_numel * 100 if non_lazy_numel != 0 else 0 + # logger = get_dist_logger() + # logger.info(f"Param lazy rate: {param_lazy_cnt}/{param_cnt}", ranks=[0]) + # logger.info(f"Buffer lazy rate: {buf_lazy_cnt}/{buf_cnt}", ranks=[0]) + # logger.info( + # f"Non lazy numel: {non_lazy_numel} ({non_lazy_numel/1024**2:.3f} M), ratio: {non_lazy_numel_ratio}%", + # ranks=[0], + # ) + + return module + + +def _is_int_tuple(args) -> bool: + if not isinstance(args, tuple): + return False + for x in args: + if not isinstance(x, int): + return False + return True + + +def _copy_tensor(tensor: Tensor, requires_grad: bool) -> Tensor: + copied = tensor.data.clone() + copied.requires_grad = requires_grad + return copied diff --git a/ixformer_sdk/train/speedformer/layers/lazy/pretrained.py b/ixformer_sdk/train/speedformer/layers/lazy/pretrained.py new file mode 100644 index 00000000..f00b18e5 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/lazy/pretrained.py @@ -0,0 +1,318 @@ +import os +from typing import Callable, Optional, Union + +import torch +from torch.nn import Module + + +class PretrainedManager: + old_from_pretrained: Optional[Callable] = None + + @staticmethod + def inject() -> None: + try: + from transformers.modeling_utils import PreTrainedModel + except ImportError: + return + # recover bound method to plain function + PretrainedManager.old_from_pretrained = PreTrainedModel.from_pretrained.__func__ + PreTrainedModel.from_pretrained = new_from_pretrained + + @staticmethod + def recover() -> None: + try: + from transformers.modeling_utils import PreTrainedModel + except ImportError: + return + # convert plain function to class method + PreTrainedModel.from_pretrained = classmethod( + PretrainedManager.old_from_pretrained) + PretrainedManager.old_from_pretrained = None + + +@classmethod +def new_from_pretrained( + cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs +) -> Module: + from transformers import GenerationConfig + from transformers.configuration_utils import PretrainedConfig + from transformers.modeling_utils import ( + ContextManagers, + _add_variant, + cached_file, + download_url, + has_file, + is_offline_mode, + is_remote_url, + no_init_weights, + ) + from transformers.utils import ( + SAFE_WEIGHTS_INDEX_NAME, + SAFE_WEIGHTS_NAME, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, + is_safetensors_available, + logging, + ) + + logger = logging.get_logger(__name__) + + config = kwargs.pop("config", None) + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + resume_download = kwargs.pop("resume_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", False) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + _ = kwargs.pop("mirror", None) + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + _fast_init = kwargs.pop("_fast_init", True) + torch_dtype = kwargs.pop("torch_dtype", None) + subfolder = kwargs.pop("subfolder", "") + commit_hash = kwargs.pop("_commit_hash", None) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop( + "use_safetensors", None if is_safetensors_available() else False) + + if len(kwargs) > 0: + logger.warning(f"Below kwargs may be ignored: {list(kwargs.keys())}") + + from_pt = True + + user_agent = {"file_type": "model", "framework": "pytorch", + "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + # Load config if we don't provide a configuration + if not isinstance(config, PretrainedConfig): + config_path = config if config is not None else pretrained_model_name_or_path + config, model_kwargs = cls.config_class.from_pretrained( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + **kwargs, + ) + else: + model_kwargs = kwargs + + if commit_hash is None: + commit_hash = getattr(config, "_commit_hash", None) + + # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the + # index of the files. + + if pretrained_model_name_or_path is not None: + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if is_local: + if use_safetensors is not False and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, + _add_variant(SAFE_WEIGHTS_NAME, variant)) + ): + # Load from a safetensors checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant( + SAFE_WEIGHTS_NAME, variant) + ) + elif use_safetensors is not False and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, + _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)) + ): + # Load from a sharded safetensors checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant( + SAFE_WEIGHTS_INDEX_NAME, variant) + ) + elif os.path.isfile( + os.path.join(pretrained_model_name_or_path, + subfolder, _add_variant(WEIGHTS_NAME, variant)) + ): + # Load from a PyTorch checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant( + WEIGHTS_NAME, variant) + ) + elif os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, + _add_variant(WEIGHTS_INDEX_NAME, variant)) + ): + # Load from a sharded PyTorch checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant( + WEIGHTS_INDEX_NAME, variant) + ) + else: + raise EnvironmentError( + f"Error no file named {_add_variant(WEIGHTS_NAME, variant)} found in directory" + f" {pretrained_model_name_or_path}." + ) + elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)): + archive_file = pretrained_model_name_or_path + is_local = True + elif is_remote_url(pretrained_model_name_or_path): + filename = pretrained_model_name_or_path + resolved_archive_file = download_url(pretrained_model_name_or_path) + else: + # set correct filename + if use_safetensors is not False: + filename = _add_variant(SAFE_WEIGHTS_NAME, variant) + else: + filename = _add_variant(WEIGHTS_NAME, variant) + + try: + # Load from URL or cache if already cached + cached_file_kwargs = { + "cache_dir": cache_dir, + "force_download": force_download, + "proxies": proxies, + "resume_download": resume_download, + "local_files_only": local_files_only, + "use_auth_token": use_auth_token, + "user_agent": user_agent, + "revision": revision, + "subfolder": subfolder, + "_raise_exceptions_for_missing_entries": False, + "_commit_hash": commit_hash, + } + resolved_archive_file = cached_file( + pretrained_model_name_or_path, filename, **cached_file_kwargs) + + # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None + # result when internet is up, the repo and revision exist, but the file does not. + if resolved_archive_file is None and filename == _add_variant(SAFE_WEIGHTS_NAME, variant): + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + resolved_archive_file = cached_file( + pretrained_model_name_or_path, + _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant), + **cached_file_kwargs, + ) + if resolved_archive_file is not None: + pass + elif use_safetensors: + raise EnvironmentError( + f" {_add_variant(SAFE_WEIGHTS_NAME, variant)} or {_add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)} and thus cannot be loaded with `safetensors`. Please make sure that the model has been saved with `safe_serialization=True` or do not set `use_safetensors=True`." + ) + else: + # This repo has no safetensors file of any kind, we switch to PyTorch. + filename = _add_variant(WEIGHTS_NAME, variant) + resolved_archive_file = cached_file( + pretrained_model_name_or_path, filename, **cached_file_kwargs + ) + if resolved_archive_file is None and filename == _add_variant(WEIGHTS_NAME, variant): + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + resolved_archive_file = cached_file( + pretrained_model_name_or_path, + _add_variant(WEIGHTS_INDEX_NAME, variant), + **cached_file_kwargs, + ) + if resolved_archive_file is not None: + pass + if resolved_archive_file is None: + # Otherwise, maybe there is a TF or Flax model file. We try those to give a helpful error + # message. + has_file_kwargs = { + "revision": revision, + "proxies": proxies, + "use_auth_token": use_auth_token, + } + if variant is not None and has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)} but there is a file without the variant" + f" {variant}. Use `variant=None` to load this model from those weights." + ) + else: + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)}" + ) + except EnvironmentError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted + # to the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + raise EnvironmentError( + f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it" + " from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a file named {_add_variant(WEIGHTS_NAME, variant)}." + ) + + if is_local: + logger.info(f"loading weights file {archive_file}") + resolved_archive_file = archive_file + else: + logger.info( + f"loading weights file {filename} from cache at {resolved_archive_file}") + else: + resolved_archive_file = None + + if from_pt: + # set dtype to instantiate the model under: + # 1. If torch_dtype is not None, we use that dtype + dtype_orig = None + + if torch_dtype is not None: + if not isinstance(torch_dtype, torch.dtype): + raise ValueError( + f"`torch_dtype` can be either `torch.dtype` or `None`, but received {torch_dtype}") + dtype_orig = cls._set_default_torch_dtype(torch_dtype) + + config.name_or_path = pretrained_model_name_or_path + + # Instantiate model. + init_contexts = [no_init_weights(_enable=_fast_init)] + + with ContextManagers(init_contexts): + model = cls(config, *model_args, **model_kwargs) + + if from_pt: + # restore default dtype + if dtype_orig is not None: + torch.set_default_dtype(dtype_orig) + + # make sure token embedding weights are still tied if needed + model.tie_weights() + + # Set model in evaluation mode to deactivate DropOut modules by default + model.eval() + + # If it is a model with generation capabilities, attempt to load the generation config + if model.can_generate(): + try: + model.generation_config = GenerationConfig.from_pretrained( + pretrained_model_name_or_path, + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + **kwargs, + ) + except (OSError, TypeError): + logger.info( + "Generation config file not found, using a generation config created from the model config.") + + return model diff --git a/ixformer_sdk/train/speedformer/layers/llama/__init__.py b/ixformer_sdk/train/speedformer/layers/llama/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/layers/llama/attention.py b/ixformer_sdk/train/speedformer/layers/llama/attention.py new file mode 100644 index 00000000..35663716 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/llama/attention.py @@ -0,0 +1,186 @@ +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +from ixformer.train.speedformer.models.llama.configuration_llama import LlamaConfig +from ixformer.train.speedformer.models.llama.modeling_llama import LlamaFlashAttention2 +from transformers import Cache +from transformers.utils import logging + +from flash_attn import flash_attn_func, flash_attn_varlen_func +from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input + +from ixformer.train.functions.fused_rope import fused_apply_rotary_pos_emb +from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding + + +class BaseLlamaAttention(LlamaFlashAttention2): + """ + 加这个层的原因：1.当原模型中使用的是torch nvtive的attention，强制替换成flash_attn; 2.优化rope + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if self.config.rope_scaling is None: + self.rotary_emb = RotaryEmbedding(self.head_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + output_attentions = False + bsz, q_len, _ = hidden_states.size() + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # fused_apply_rotary_pos_emb need qk to be in "sbhd" + query_states = query_states.view( + bsz, q_len, self.num_heads, self.head_dim).transpose(1, 0).contiguous() + key_states = key_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 0).contiguous() + value_states = value_states.view( + bsz, q_len, self.num_heads, self.head_dim) + + kv_seq_len = key_states.shape[0] + if past_key_value is not None: + kv_seq_len += past_key_value[0].shape[0] + + emb = self.rotary_emb(kv_seq_len).to(dtype=torch.float32) + query_states = fused_apply_rotary_pos_emb(query_states, emb) + key_states = fused_apply_rotary_pos_emb(key_states, emb) + + # kv cache staff + if past_key_value is not None: + # reuse k, v, self_attention + key_states = torch.cat([past_key_value[0], key_states], dim=0) + value_states = torch.cat([past_key_value[1], value_states], dim=0) + past_key_value = (key_states, value_states) if use_cache else None + + dropout_rate = self.attention_dropout if self.training else 0.0 + + # after fused_apply_rotary_pos_emb, qk change to "bshd" for flashattn or "bhsd" for sdpa + if attention_mask is None: # flash-attn + query_states = query_states.transpose(0, 1).contiguous() + key_states = key_states.transpose(0, 1).contiguous() + else: # sdpa + query_states = query_states.permute(1, 2, 0, 3).contiguous() + key_states = key_states.permute(1, 2, 0, 3).contiguous() + value_states = value_states.transpose(1, 2).contiguous() + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (LlamaRMSNorm handles it correctly) + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + # Handle the case where the model is quantized + if hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = self._flash_attention_forward( + query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate + ) + + attn_output = attn_output.reshape( + bsz, q_len, self.hidden_size).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def _flash_attention_forward( + self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None + ): + """ + for now, if attention_mask is none, flash-attn has better performance than torch.nn.functional.scaled_dot_product_attention; + if attention_mask is not none, torch.nn.functional.scaled_dot_product_attention works better + so sdpa and flash-attn is perfered according to attention_mask + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`int`, *optional*): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + """ + # Contains at least one padding token in the sequence + # if attention_mask is not None: + if attention_mask is not None: + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attention_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal=self.is_causal and attention_mask is None and query_length > 1, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + + else: + attn_output = flash_attn_func( + query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=self.is_causal + ) + + return attn_output + + +class LlamaAttention(BaseLlamaAttention): + def __init__(self) -> None: + raise NotImplementedError( + "LlamaAttention is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to LlamaAttention module provided above." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + + # LazyInitContext.materialize(module) + + # try to get normalized_shape, eps, elementwise_affine from the module + config = getattr(module, "config") + layer_idx = getattr(module, "layer_idx", None) + + attention = BaseLlamaAttention( + config=config, + layer_idx=layer_idx, + ) + + attention.q_proj.weight = module.q_proj.weight + attention.k_proj.weight = module.k_proj.weight + attention.v_proj.weight = module.v_proj.weight + attention.o_proj.weight = module.o_proj.weight + + if config.attention_bias: + attention.q_proj.bias = module.q_proj.bias + attention.k_proj.bias = module.k_proj.bias + attention.v_proj.bias = module.v_proj.bias + attention.o_proj.bias = module.o_proj.bias + return attention diff --git a/ixformer_sdk/train/speedformer/layers/llama/llama_method.py b/ixformer_sdk/train/speedformer/layers/llama/llama_method.py new file mode 100644 index 00000000..274b55a6 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/llama/llama_method.py @@ -0,0 +1,224 @@ +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ixformer.train.speedformer.models.llama.modeling_llama import LlamaModel +from ixformer.train.speedformer.models.llama.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask_for_sdpa +from ixformer.train.speedformer.layers.cross_entropy_loss import fast_cross_entropy_loss as CrossEntropyLoss +from transformers.utils import logging +from transformers.cache_utils import Cache, DynamicCache + +logger = logging.get_logger(__name__) + + +def LlamaModel_forward(): + from transformers.modeling_outputs import BaseModelOutputWithPast + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError( + "You have to specify either input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + past_key_values_length = 0 + if use_cache: + use_legacy_cache = not isinstance(past_key_values, Cache) + if use_legacy_cache: + past_key_values = DynamicCache.from_legacy_cache( + past_key_values) + past_key_values_length = past_key_values.get_usable_length( + seq_length) + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if attention_mask is not None: + # output_attentions=True can not be supported when using SDPA, and we fall back on + # the manual implementation that requires a 4D causal mask in all cases. + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + ) + + # embed positions + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache.to_legacy_cache( + ) if use_legacy_cache else next_decoder_cache + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + return forward + + +def LlamaForCausalLM_forward(): + from transformers.utils import add_start_docstrings_to_model_forward, replace_return_docstrings + from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING, CausalLMOutputWithPast, _CONFIG_FOR_DOC + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + if self.config.pretraining_tp > 1: + lm_head_slices = self.lm_head.weight.split( + self.vocab_size // self.config.pretraining_tp, dim=0) + logits = [F.linear(hidden_states, lm_head_slices[i]) + for i in range(self.config.pretraining_tp)] + logits = torch.cat(logits, dim=-1) + else: + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + return forward diff --git a/ixformer_sdk/train/speedformer/layers/llama/mlp.py b/ixformer_sdk/train/speedformer/layers/llama/mlp.py new file mode 100644 index 00000000..9c7ace2d --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/llama/mlp.py @@ -0,0 +1,55 @@ +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +import ixformer.train.functions as F +from ixformer.train.speedformer.models.llama.configuration_llama import LlamaConfig +from ixformer.train.speedformer.models.llama.modeling_llama import LlamaMLP +from transformers import Cache +from transformers.utils import logging + +from ixformer.train.speedformer.layers.lazy import LazyInitContext + + +class BaseLlamaMLP(LlamaMLP): + """ + 这个层主要的优化点是：将linear1(act(cat(linear2(x), linear3(x))))的结构变成 linear1(act(linear23(x))) + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.gate_up = nn.Linear( + self.hidden_size, self.intermediate_size * 2, bias=False) + del self.gate_proj, self.up_proj + del self.act_fn + + def forward(self, x): + res = self.gate_up(x) + down_proj = self.down_proj(F.swiglu(res)) + return down_proj + + +class IXFLlamaMLP(BaseLlamaMLP): + def __init__(self) -> None: + raise NotImplementedError( + "IXFLlamaMLP is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to IXFLlamaMLP module provided above." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + + LazyInitContext.materialize(module) + + config = getattr(module, "config") + + mlp = BaseLlamaMLP(config=config) + + mlp.gate_up.weight.data = torch.concat( + (module.gate_proj.weight.data, module.up_proj.weight.data), dim=0) + mlp.down_proj.weight.data = module.down_proj.weight.data + + return mlp diff --git a/ixformer_sdk/train/speedformer/layers/normalization.py b/ixformer_sdk/train/speedformer/layers/normalization.py new file mode 100644 index 00000000..b3c025d0 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/normalization.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- +import warnings +from abc import ABC, abstractmethod + +import torch +import torch.nn as nn +import ixformer.functions as ixff +from ixformer.train.functions import FusedRMSNorm as ixf_FusedRMSNorm +from apex.normalization.fused_layer_norm import FusedRMSNorm as apex_FusedRMSNorm +from ixformer.train.speedformer.layers.lazy import LazyInitContext + + +class BaseLayerNorm(ABC): + @abstractmethod + def from_native_module(module: nn.Module, sp_partial_derived: bool = False): + """ + Convert a native PyTorch layer normalization module to a specific layer normalization module, + and optionally mark parameters for gradient aggregation. + + Args: + module (nn.Module): The native PyTorch layer normalization module to be converted. + sp_partial_derived (bool): Whether this module's gradients are partially derived in sequence parallelism. + + Returns: + nn.Module: The specific layer normalization module. + + Raises: + AssertionError: If the provided module is not an instance of the supported layer normalization type. + """ + + +class IXFFusedRMSNorm(BaseLayerNorm): + """ + This is a wrapper around the apex fused rms norm implementation. It is meant to be used only with the from_native_module interface. + """ + + def __init__(self) -> None: + raise NotImplementedError( + "FusedRMSNorm is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native RMSNorm module to FusedRMSNorm module provided by apex." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + r""" + Convert a native RMSNorm module module to FusedRMSNorm module provided by ixformer, + and optionally marking parameters for gradient aggregation. + + Args: + module (nn.LayerNorm): The native PyTorch LayerNorm module to be converted. + sp_partial_derived (bool): Whether this module's gradients are partially derived in sequence parallelism. + + Returns: + nn.Module: FusedRMSNorm module. + """ + + LazyInitContext.materialize(module) + + # try to get normalized_shape, eps, elementwise_affine from the module + normalized_shape = getattr( + module, "normalized_shape", module.weight.shape[0]) + eps = module.variance_epsilon if hasattr( + module, "variance_epsilon") else module.eps + elementwise_affine = getattr(module, "elementwise_affine", True) + + rmsnorm = ixf_FusedRMSNorm( + normalized_shape=normalized_shape, + eps=eps, + elementwise_affine=elementwise_affine, + ) + + rmsnorm.weight = module.weight + + return rmsnorm + + +class APEXFusedRMSNorm(BaseLayerNorm): + """ + This is a wrapper around the apex fused rms norm implementation. It is meant to be used only with the from_native_module interface. + """ + + def __init__(self) -> None: + raise NotImplementedError( + "FusedRMSNorm is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native RMSNorm module to FusedRMSNorm module provided by apex." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + r""" + Convert a native RMSNorm module module to FusedRMSNorm module provided by ixformer, + and optionally marking parameters for gradient aggregation. + + Args: + module (nn.LayerNorm): The native PyTorch LayerNorm module to be converted. + sp_partial_derived (bool): Whether this module's gradients are partially derived in sequence parallelism. + + Returns: + nn.Module: FusedRMSNorm module. + """ + + LazyInitContext.materialize(module) + + # try to get normalized_shape, eps, elementwise_affine from the module + normalized_shape = getattr( + module, "normalized_shape", module.weight.shape[0]) + eps = module.variance_epsilon if hasattr( + module, "variance_epsilon") else module.eps + elementwise_affine = getattr(module, "elementwise_affine", True) + + rmsnorm = apex_FusedRMSNorm( + normalized_shape=normalized_shape, + eps=eps, + elementwise_affine=elementwise_affine, + ) + + rmsnorm.weight = module.weight + + return rmsnorm + + +# 替换torch LayerNorm 的forward +@staticmethod +def replace_layernorm_forward(self, input: torch.Tensor) -> torch.Tensor: + + output = torch.empty_like(input) + + return ixff.layernorm_train(input, self.weight, self.bias, self.normalized_shape, output, True) diff --git a/ixformer_sdk/train/speedformer/layers/qwen2/__init__.py b/ixformer_sdk/train/speedformer/layers/qwen2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/layers/qwen2/attention.py b/ixformer_sdk/train/speedformer/layers/qwen2/attention.py new file mode 100644 index 00000000..be3b9b8b --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/qwen2/attention.py @@ -0,0 +1,263 @@ +import math +import warnings +import inspect +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +from ixformer.train.speedformer.models.qwen2.configuration_qwen2 import Qwen2Config +from ixformer.train.speedformer.models.qwen2.modeling_qwen2 import Qwen2FlashAttention2 +from transformers import Cache +from transformers.utils import logging + +from flash_attn import flash_attn_func, flash_attn_varlen_func + +from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input + +from ixformer.train.functions.fused_rope import fused_apply_rotary_pos_emb +from ixformer.train.speedformer.layers.rotary_pos_embedding import RotaryEmbedding + +from ixformer.train.speedformer.layers.lazy import LazyInitContext + +_flash_supports_window_size = "window_size" in list( + inspect.signature(flash_attn_func).parameters) +logger = logging.get_logger(__name__) + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class BaseQwenAttention(Qwen2FlashAttention2): + """ + 加这个层的原因：1.当原模型中使用的是torch nvtive的attention，强制替换成flash_attn; 2.优化rope + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + out_dim = self.num_heads * self.head_dim + \ + self.num_key_value_heads * self.head_dim * 2 + self.qkv_proj = nn.Linear(self.hidden_size, out_dim, bias=True) + del self.q_proj, self.k_proj, self.v_proj + self.rotary_emb = RotaryEmbedding(self.head_dim, self.rope_theta) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + ): + bsz, q_len, _ = hidden_states.size() + qkv = self.qkv_proj(hidden_states) + q_dim = self.num_heads * self.head_dim + kv_dim = self.num_key_value_heads * self.head_dim + query_states, key_states, value_states = torch.split( + qkv, (q_dim, kv_dim, kv_dim), dim=-1) + # fused_apply_rotary_pos_emb need qk to be in "sbhd", v stay "bshd" + query_states = query_states.view( + bsz, q_len, self.num_heads, self.head_dim).transpose(0, 1).contiguous() + key_states = key_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(0, 1).contiguous() + value_states = value_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim) + + kv_seq_len = key_states.shape[0] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value[0].shape[0] + + emb = self.rotary_emb(kv_seq_len).to(dtype=torch.float32) + query_states = fused_apply_rotary_pos_emb(query_states, emb) + key_states = fused_apply_rotary_pos_emb(key_states, emb) + use_sliding_windows = ( + _flash_supports_window_size + and getattr(self.config, "sliding_window", None) is not None + and kv_seq_len > self.config.sliding_window + and self.config.use_sliding_window + ) + + if not _flash_supports_window_size: + logger.warning_once( + "The current flash attention version does not support sliding window attention, for a more memory efficient implementation" + " make sure to upgrade flash-attn library." + ) + + # for now, attention with sliding_windows have not test, so if use_sliding_windows throw error + if use_sliding_windows: + raise KeyError("use_sliding_windows not support for now") + + # kv cache staff + if past_key_value is not None: + # reuse k, v, self_attention + key_states = torch.cat([past_key_value[0], key_states], dim=0) + value_states = torch.cat([past_key_value[1], value_states], dim=0) + past_key_value = (key_states, value_states) if use_cache else None + + # if attention mask is None, use flashattn which support GQA + if attention_mask is not None: + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + dropout_rate = 0.0 if not self.training else self.attention_dropout + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # after fused_apply_rotary_pos_emb, qk change to "bshd" for flashattn or "bhsd" for sdpa + if attention_mask is None: # flash-attn + query_states = query_states.transpose(0, 1).contiguous() + key_states = key_states.transpose(0, 1).contiguous() + else: # sdpa + query_states = query_states.permute(1, 2, 0, 3).contiguous() + key_states = key_states.permute(1, 2, 0, 3).contiguous() + value_states = value_states.transpose(1, 2).contiguous() + + attn_output = self._attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + dropout=dropout_rate, + use_sliding_windows=use_sliding_windows, + ) + + attn_output = attn_output.reshape( + bsz, q_len, self.hidden_size).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def _attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + query_length, + dropout=0.0, + softmax_scale=None, + use_sliding_windows=False, + ): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`float`): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + use_sliding_windows (`bool`, *optional*): + Whether to activate sliding window attention. + """ + if not self._flash_attn_uses_top_left_mask: + causal = self.is_causal + else: + # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__. + causal = self.is_causal and query_length != 1 + + if attention_mask is not None: + batch_size = query_states.shape[0] + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attention_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal=causal, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + else: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + + return attn_output + + +class QwenAttention(BaseQwenAttention): + def __init__(self) -> None: + raise NotImplementedError( + "LlamaAttention is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native LlamaAttention module to LlamaAttention module provided above." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + + LazyInitContext.materialize(module) + + # try to get normalized_shape, eps, elementwise_affine from the module + config = getattr(module, "config") + layer_idx = getattr(module, "layer_idx", None) + + attention = BaseQwenAttention( + config=config, + layer_idx=layer_idx, + ) + + attention.qkv_proj.weight.data = torch.cat( + (module.q_proj.weight.data, module.k_proj.weight.data, module.v_proj.weight.data), dim=0) + attention.qkv_proj.bias.data = torch.cat( + (module.q_proj.bias.data, module.k_proj.bias.data, module.v_proj.bias.data), dim=0) + + attention.o_proj.weight.data = module.o_proj.weight.data + + return attention diff --git a/ixformer_sdk/train/speedformer/layers/qwen2/mlp.py b/ixformer_sdk/train/speedformer/layers/qwen2/mlp.py new file mode 100644 index 00000000..d6dfaa19 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/qwen2/mlp.py @@ -0,0 +1,55 @@ +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +import ixformer.train.functions as F +from ixformer.train.speedformer.models.qwen2.configuration_qwen2 import Qwen2Config +from ixformer.train.speedformer.models.qwen2.modeling_qwen2 import Qwen2MLP +from transformers import Cache +from transformers.utils import logging + +from ixformer.train.speedformer.layers.lazy import LazyInitContext + + +class BaseQwen2MLP(Qwen2MLP): + """ + 这个层主要的优化点是：将linear1(act(cat(linear2(x), linear3(x))))的结构变成 linear1(act(linear23(x))) + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.gate_up = nn.Linear( + self.hidden_size, self.intermediate_size * 2, bias=False) + del self.gate_proj, self.up_proj + del self.act_fn + + def forward(self, x): + res = self.gate_up(x) + down_proj = self.down_proj(F.swiglu(res)) + return down_proj + + +class IXFQwen2MLP(BaseQwen2MLP): + def __init__(self) -> None: + raise NotImplementedError( + "IXFQwen2MLP is not implemented as a physical class. " + "It is meant to be used only with the from_native_module interface to Convert a native Qwen2MLP module to BaseQwen2MLP module provided above." + ) + + @staticmethod + def from_native_module(module: nn.Module, *args, **kwargs) -> nn.Module: + + LazyInitContext.materialize(module) + + config = getattr(module, "config") + + mlp = BaseQwen2MLP(config=config) + + mlp.gate_up.weight.data = torch.concat( + (module.gate_proj.weight.data, module.up_proj.weight.data), dim=0) + mlp.down_proj.weight.data = module.down_proj.weight.data + + return mlp diff --git a/ixformer_sdk/train/speedformer/layers/rotary_pos_embedding.py b/ixformer_sdk/train/speedformer/layers/rotary_pos_embedding.py new file mode 100644 index 00000000..18af0501 --- /dev/null +++ b/ixformer_sdk/train/speedformer/layers/rotary_pos_embedding.py @@ -0,0 +1,55 @@ +import importlib.util +import torch + +from torch import einsum, nn + +__all__ = ['RotaryEmbedding'] + + +# RotaryEmbedding and apply_rotary_pos_emb are copy from http://bitbucket.iluvatar.ai:7990/projects/PSR/repos/megatron-deepspeed/browse/megatron/model/rotary_pos_embedding.py +# for now RotaryEmbedding is used, apply_rotary_pos_emb can be replaced by fused_apply_rotary_pos_emb from ixformer for better performance + +class RotaryEmbedding(nn.Module): + def __init__(self, dim, base=10000): + super().__init__() + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) + self.register_buffer('inv_freq', inv_freq) + if importlib.util.find_spec('einops') is None: + raise RuntimeError("einops is required for Rotary Embedding") + + def forward(self, max_seq_len, offset=0): + seq = torch.arange(max_seq_len, device=self.inv_freq.device) + offset + freqs = einsum( + 'i , j -> i j', seq.type_as(self.inv_freq), self.inv_freq) + # first part even vector components, second part odd vector components, + # 2 * dim in dimension size + emb = torch.cat((freqs, freqs), dim=-1) + # emb [seq_length, .., dim] + from einops import rearrange + return rearrange(emb, 'n d -> n 1 1 d') + + +def _rotate_half(x): + """ + change sign so the last dimension becomes [-odd, +even] + """ + from einops import rearrange + x = rearrange(x, '... (j d) -> ... j d', j=2) + x1, x2 = x.unbind(dim=-2) + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(t, freqs): + """ + input tensor t is of shape [seq_length, ..., dim] + rotary positional embeding tensor freqs is of shape [seq_length, ..., dim] + check https://kexue.fm/archives/8265 for detailed formulas + """ + rot_dim = freqs.shape[-1] + # ideally t_pass is empty so rotary pos embedding is applied to all tensor t + t, t_pass = t[..., :rot_dim], t[..., rot_dim:] + + # first part is cosine component + # second part is sine component, need to change signs with _rotate_half method + t = (t * freqs.cos()) + (_rotate_half(t) * freqs.sin()) + return torch.cat((t, t_pass), dim=-1) diff --git a/ixformer_sdk/train/speedformer/model_replacer_mapping.py b/ixformer_sdk/train/speedformer/model_replacer_mapping.py new file mode 100644 index 00000000..f7df38d3 --- /dev/null +++ b/ixformer_sdk/train/speedformer/model_replacer_mapping.py @@ -0,0 +1,18 @@ +import torch + +from ixformer.train.speedformer.policy.gpt2 import GPT2Replacer +from ixformer.train.speedformer.policy.qwen2 import Qwen2Replacer +from ixformer.train.speedformer.policy.llama import LlamaReplacer +from ixformer.train.speedformer.policy.baichuan import BaichuanReplacer +from ixformer.train.speedformer.policy.bloom import BloomReplacer +from ixformer.train.speedformer.policy.chatglm import ChatglmReplacer + + +ModelMapping = { + "gpt2": GPT2Replacer, + "qwen2": Qwen2Replacer, + "llama": LlamaReplacer, + "baichuan": BaichuanReplacer, + "bloom": BloomReplacer, + "chatglm": ChatglmReplacer +} diff --git a/ixformer_sdk/train/speedformer/models/__init__.py b/ixformer_sdk/train/speedformer/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/models/baichuan/__init__.py b/ixformer_sdk/train/speedformer/models/baichuan/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/models/baichuan/configuration_baichuan.py b/ixformer_sdk/train/speedformer/models/baichuan/configuration_baichuan.py new file mode 100644 index 00000000..e067bb7a --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/baichuan/configuration_baichuan.py @@ -0,0 +1,68 @@ +# Copyright 2023 Baichuan Inc. All Rights Reserved. + +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + + +logger = logging.get_logger(__name__) + + +class BaichuanConfig(PretrainedConfig): + model_type = "baichuan" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=125696, + hidden_size=4096, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + hidden_act="silu", + max_position_embeddings=4096, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + tie_word_embeddings=False, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.z_loss_weight = 0 + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/ixformer_sdk/train/speedformer/models/baichuan/generation_utils.py b/ixformer_sdk/train/speedformer/models/baichuan/generation_utils.py new file mode 100644 index 00000000..57716991 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/baichuan/generation_utils.py @@ -0,0 +1,83 @@ +from typing import List +from queue import Queue + +import torch + + +def build_chat_input(model, tokenizer, messages: List[dict], max_new_tokens: int=0): + def _parse_messages(messages, split_role="user"): + system, rounds = "", [] + round = [] + for i, message in enumerate(messages): + if message["role"] == "system": + assert i == 0 + system = message["content"] + continue + if message["role"] == split_role and round: + rounds.append(round) + round = [] + round.append(message) + if round: + rounds.append(round) + return system, rounds + + max_new_tokens = max_new_tokens or model.generation_config.max_new_tokens + max_input_tokens = model.config.model_max_length - max_new_tokens + system, rounds = _parse_messages(messages, split_role="user") + system_tokens = tokenizer.encode(system) + max_history_tokens = max_input_tokens - len(system_tokens) + + history_tokens = [] + for round in rounds[::-1]: + round_tokens = [] + for message in round: + if message["role"] == "user": + round_tokens.append(model.generation_config.user_token_id) + else: + round_tokens.append(model.generation_config.assistant_token_id) + round_tokens.extend(tokenizer.encode(message["content"])) + if len(history_tokens) == 0 or len(history_tokens) + len(round_tokens) <= max_history_tokens: + history_tokens = round_tokens + history_tokens # concat left + if len(history_tokens) < max_history_tokens: + continue + break + + input_tokens = system_tokens + history_tokens + if messages[-1]["role"] != "assistant": + input_tokens.append(model.generation_config.assistant_token_id) + input_tokens = input_tokens[-max_input_tokens:] # truncate left + return torch.LongTensor([input_tokens]).to(model.device) + + +class TextIterStreamer: + def __init__(self, tokenizer, skip_prompt=False, skip_special_tokens=False): + self.tokenizer = tokenizer + self.skip_prompt = skip_prompt + self.skip_special_tokens = skip_special_tokens + self.tokens = [] + self.text_queue = Queue() + self.next_tokens_are_prompt = True + + def put(self, value): + if self.skip_prompt and self.next_tokens_are_prompt: + self.next_tokens_are_prompt = False + else: + if len(value.shape) > 1: + value = value[0] + self.tokens.extend(value.tolist()) + self.text_queue.put( + self.tokenizer.decode(self.tokens, skip_special_tokens=self.skip_special_tokens)) + + def end(self): + self.text_queue.put(None) + + def __iter__(self): + return self + + def __next__(self): + value = self.text_queue.get() + if value is None: + raise StopIteration() + else: + return value + diff --git a/ixformer_sdk/train/speedformer/models/baichuan/modeling_baichuan.py b/ixformer_sdk/train/speedformer/models/baichuan/modeling_baichuan.py new file mode 100644 index 00000000..cc0c8804 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/baichuan/modeling_baichuan.py @@ -0,0 +1,783 @@ +# Copyright 2023 Baichuan Inc. All Rights Reserved. + +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. + + +from .configuration_baichuan import BaichuanConfig +from .generation_utils import build_chat_input, TextIterStreamer + +import math +from typing import List, Optional, Tuple, Union +from threading import Thread + +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from torch.nn import functional as F +from transformers import PreTrainedModel, PretrainedConfig +from transformers.activations import ACT2FN +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.generation.utils import GenerationConfig +from transformers.utils import logging, ContextManagers + +import os +from contextlib import contextmanager +logger = logging.get_logger(__name__) + +try: + from xformers import ops as xops +except ImportError: + xops = None + logger.warning( + "Xformers is not installed correctly. If you want to use memory_efficient_attention to accelerate training use the following command to install Xformers\npip install xformers." + ) + + +# Copied from transformers.models.bart.modeling_bart._make_causal_mask +def _make_causal_mask( + input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0 +): + """ + Make causal mask used for bi-directional self-attention. + """ + bsz, tgt_len = input_ids_shape + mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device) + mask_cond = torch.arange(mask.size(-1), device=device) + mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0) + mask = mask.to(dtype) + + if past_key_values_length > 0: + mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1) + return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) + +def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + if len(mask.size()) == 3: + bsz, src_len, _ = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + expanded_mask = mask[:,None,:,:].expand(bsz, 1, tgt_len, src_len).to(dtype) + else: + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) + + +class RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + + # convert into half-precision if necessary + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + + return self.weight * hidden_states + + +class RotaryEmbedding(torch.nn.Module): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + super().__init__() + self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim)) + self.max_seq_len_cached = max_position_embeddings + t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32) + freqs = torch.outer(t, self.inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32) + self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32) + def forward(self, x, seq_len=None): + # x: [bs, num_attention_heads, seq_len, head_size] + # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case. + if seq_len > self.max_seq_len_cached: + self.max_seq_len_cached = seq_len + t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32) + freqs = torch.outer(t, self.inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32).to(x.device) + self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32).to(x.device) + elif self.cos_cached.device != x.device: + self.cos_cached = self.cos_cached.to(x.device) + self.sin_cached = self.sin_cached.to(x.device) + return ( + self.cos_cached[:, :, :seq_len, ...], + self.sin_cached[:, :, :seq_len, ...], + ) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2:] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos_, sin_, position_ids): + cos = cos_.squeeze(1).squeeze(0) # [seq_len, dim] + sin = sin_.squeeze(1).squeeze(0) # [seq_len, dim] + cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] + sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] + q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin) + k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin) + return q_embed.to(q.dtype), k_embed.to(k.dtype) + + +class MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + ): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.act_fn = ACT2FN[hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + def __init__(self, config: BaichuanConfig): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.max_position_embeddings = config.max_position_embeddings + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + self.W_pack = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + self.rotary_emb = RotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: bool = False, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + proj = self.W_pack(hidden_states) + proj = proj.unflatten(-1, (3, self.hidden_size)).unsqueeze(0).transpose(0, -2).squeeze(-2) + query_states = proj[0].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = proj[1].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + value_states = proj[2].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value[0].shape[-2] + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + # [bsz, nh, t, hd] + + if past_key_value is not None: + # reuse k, v, self_attention + key_states = torch.cat([past_key_value[0], key_states], dim=2) + value_states = torch.cat([past_key_value[1], value_states], dim=2) + + past_key_value = (key_states, value_states) if use_cache else None + if xops is not None and self.training: + attn_weights = None + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + attn_output = xops.memory_efficient_attention( + query_states, key_states, value_states, attn_bias=xops.LowerTriangularMask() + ) + else: + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True): + attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask = attention_mask) + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class DecoderLayer(nn.Module): + def __init__(self, config: BaichuanConfig): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Attention(config=config) + self.mlp = MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +class BaichuanPreTrainedModel(PreTrainedModel): + config_class = BaichuanConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["DecoderLayer"] + _keys_to_ignore_on_load_unexpected = [r"decoder\.version"] + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, BaichuanModel): + module.gradient_checkpointing = value + + +class BaichuanModel(BaichuanPreTrainedModel): + def __init__(self, config: BaichuanConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask + def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length): + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + combined_attention_mask = None + if input_shape[-1] > 1: + combined_attention_mask = _make_causal_mask( + input_shape, + inputs_embeds.dtype, + device=inputs_embeds.device, + past_key_values_length=past_key_values_length, + ) + + if attention_mask is not None: + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to( + inputs_embeds.device + ) + combined_attention_mask = ( + expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask + ) + + return combined_attention_mask + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + seq_length_with_past = seq_length + past_key_values_length = 0 + + if past_key_values is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + # embed positions + if attention_mask is None: + attention_mask = torch.ones( + (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device + ) + attention_mask = self._prepare_decoder_attention_mask( + attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length + ) + + hidden_states = inputs_embeds + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = () if use_cache else None + + for idx, decoder_layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + past_key_value = past_key_values[idx] if past_key_values is not None else None + + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + # None for past_key_value + return module(*inputs, output_attentions, None) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(decoder_layer), + hidden_states, + attention_mask, + position_ids, + None, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class NormHead(nn.Module): + def __init__(self, hidden_size, vocab_size, bias=False): + super().__init__() + self.weight = nn.Parameter(torch.empty((vocab_size, hidden_size))) + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + self.first_flag = True + + def forward(self, hidden_states): + if self.training: + norm_weight = nn.functional.normalize(self.weight) + elif self.first_flag: + self.first_flag = False + self.weight = nn.Parameter(nn.functional.normalize(self.weight)) + norm_weight = self.weight + else: + norm_weight = self.weight + return nn.functional.linear(hidden_states, norm_weight) + +_init_weights = True +@contextmanager +def no_init_weights(_enable=True): + global _init_weights + old_init_weights = _init_weights + if _enable: + _init_weights = False + try: + yield + finally: + _init_weights = old_init_weights + +class BaichuanForCausalLM(BaichuanPreTrainedModel): + def __init__(self, config, *model_args, **model_kwargs): + super().__init__(config, *model_args, **model_kwargs) + self.model = BaichuanModel(config) + + self.lm_head = NormHead(config.hidden_size, config.vocab_size, bias=False) + if hasattr(config, "quantization_config") and config.quantization_config['load_in_4bit']: + try: + from .quantizer import quantize_offline, init_model_weight_int4 + except ImportError: + raise ImportError(f"Needs QLinear to run quantize.") + quantize_offline(self, 4) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + *model_args, + config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None, + cache_dir: Optional[Union[str, os.PathLike]] = None, + ignore_mismatched_sizes: bool = False, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + use_safetensors: bool = None, + **kwargs, + ): + # Load config if we don't provide a configuration + if not isinstance(config, PretrainedConfig): + config_path = config if config is not None else pretrained_model_name_or_path + config, model_kwargs = cls.config_class.from_pretrained( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + force_download=force_download, + resume_download=False, + proxies=None, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder="", + _from_auto=False, + _from_pipeline=None, + **kwargs, + ) + else: + model_kwargs = kwargs + + if hasattr(config, "quantization_config") and config.quantization_config['load_in_4bit']: + try: + from .quantizer import init_model_weight_int4 + from accelerate import init_empty_weights, dispatch_model, infer_auto_device_map + from accelerate.utils import CustomDtype + from accelerate.utils import get_balanced_memory + except ImportError: + raise ImportError(f"Needs import model weight init func to run quantize.") + # Instantiate model. + init_contexts = [no_init_weights(_enable=True)] + init_contexts.append(init_empty_weights()) + with ContextManagers(init_contexts): + model = cls(config) + + model_file = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin') + state_dict = torch.load(model_file, map_location="cpu") + model.is_quantized = True + + device_map = kwargs.pop("device_map", None) + torch_dtype = kwargs.pop("torch_dtype", None) + + kwargs = {"no_split_module_classes": model._no_split_modules} + target_dtype = CustomDtype.INT4 + max_memory = get_balanced_memory( + model, + dtype=target_dtype, + low_zero=(device_map == "balanced_low_0"), + max_memory=None, + **kwargs, + ) + kwargs["max_memory"] = max_memory + + device_map = infer_auto_device_map(model, dtype=target_dtype, **kwargs) + model = init_model_weight_int4(config, model, state_dict) + + # Set model in evaluation mode to deactivate DropOut modules by default + model.eval() + # If it is a model with generation capabilities, attempt to load the generation config + if model.can_generate(): + try: + model.generation_config = GenerationConfig.from_pretrained( + pretrained_model_name_or_path, + cache_dir=cache_dir, + force_download=force_download, + resume_download=False, + proxies=None, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder="", + _from_auto=False, + _from_pipeline=None, + **kwargs, + ) + except (OSError, TypeError): + logger.info( + "Generation config file not found, using a generation config created from the model config." + ) + pass + + if device_map is not None: + dispatch_model(model, device_map=device_map) + + return model + return super(BaichuanForCausalLM, cls).from_pretrained(pretrained_model_name_or_path, *model_args, + config=config, cache_dir=cache_dir, ignore_mismatched_sizes=ignore_mismatched_sizes, + force_download=force_download, local_files_only=local_files_only, token=token, revision=revision, + use_safetensors=use_safetensors, **kwargs) + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + softmax_normalizer = shift_logits.max(-1).values ** 2 + z_loss = self.config.z_loss_weight * softmax_normalizer.mean() + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + z_loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs + ): + if past_key_values: + input_ids = input_ids[:, -1:] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -1].unsqueeze(-1) + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + } + ) + return model_inputs + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),) + return reordered_past + + def quantize(self, bits: int): + try: + from .quantizer import quantize_online + except ImportError: + raise ImportError(f"Needs QLinear to run quantize.") + return quantize_online(self, bits) + + def chat(self, tokenizer, messages: List[dict], stream=False, + generation_config: Optional[GenerationConfig]=None): + generation_config = generation_config or self.generation_config + input_ids = build_chat_input(self, tokenizer, messages, generation_config.max_new_tokens) + if stream: + streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) + Thread(target=self.generate, kwargs=dict( + inputs=input_ids, streamer=streamer, + generation_config=generation_config, + )).start() + return streamer + else: + outputs = self.generate(input_ids, generation_config=generation_config) + response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True) + return response \ No newline at end of file diff --git a/ixformer_sdk/train/speedformer/models/baichuan/quantizer.py b/ixformer_sdk/train/speedformer/models/baichuan/quantizer.py new file mode 100644 index 00000000..239a2fbf --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/baichuan/quantizer.py @@ -0,0 +1,210 @@ +import bitsandbytes as bnb +from bitsandbytes.nn.modules import Params4bit, Int8Params +import torch + +def Params4bitCuda(self, device): + self.data = self.data.cuda(device) + self.quant_state[0] = self.quant_state[0].cuda(device) + self.quant_state[4][0] = self.quant_state[4][0].cuda(device) + self.quant_state[4][1][0] = self.quant_state[4][1][0].cuda(device) + self.quant_state[4][1][1] = self.quant_state[4][1][1].cuda(device) + + self.quant_state[6] = self.quant_state[6].cuda(device) + return self + +class Linear4bitOnline(torch.nn.Module): + def __init__(self, weight, bias, quant_type): + super().__init__() + self.weight = Params4bit( + weight.data, requires_grad=False, compress_statistics=True, quant_type=quant_type + ) + self.compute_dtype = None + #self.weight.cuda(weight.device) + self.bias = bias + + def forward(self, x: torch.Tensor): + # weights are cast automatically as Int8Params, but the bias has to be cast manually + if self.bias is not None and self.bias.dtype != x.dtype: + self.bias.data = self.bias.data.to(x.dtype) + + if getattr(self.weight, "quant_state", None) is None: + print( + "FP4 quantization state not initialized. Please call .cuda() or .to(device) on the LinearFP4 layer first." + ) + inp_dtype = x.dtype + if self.compute_dtype is not None: + x = x.to(self.compute_dtype) + + bias = None if self.bias is None else self.bias.to(self.compute_dtype) + out = bnb.matmul_4bit( + x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state + ) + + out = out.to(inp_dtype) + + return out + +class Linear8bitLtOnline(torch.nn.Module): + def __init__( + self, + weight, + bias, + has_fp16_weights=True, + memory_efficient_backward=False, + threshold=0.0, + index=None, + ): + super().__init__() + assert ( + not memory_efficient_backward + ), "memory_efficient_backward is no longer required and the argument is deprecated in 0.37.0 and will be removed in 0.39.0" + self.state = bnb.MatmulLtState() + self.index = index + + # Necessary for stacked layers + self.state.threshold = threshold + self.state.has_fp16_weights = has_fp16_weights + self.state.memory_efficient_backward = memory_efficient_backward + if threshold > 0.0 and not has_fp16_weights: + self.state.use_pool = True + + self.weight = Int8Params( + weight.data, + has_fp16_weights=has_fp16_weights, + requires_grad=has_fp16_weights, + ) + self.bias = bias + + def init_8bit_state(self): + self.state.CB = self.weight.CB + self.state.SCB = self.weight.SCB + self.weight.CB = None + self.weight.SCB = None + + def forward(self, x: torch.Tensor): + self.state.is_training = self.training + if self.weight.CB is not None: + self.init_8bit_state() + + # weights are cast automatically as Int8Params, but the bias has to be cast manually + if self.bias is not None and self.bias.dtype != x.dtype: + self.bias.data = self.bias.data.to(x.dtype) + + out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state) + + if not self.state.has_fp16_weights: + if self.state.CB is not None and self.state.CxB is not None: + # we converted 8-bit row major to turing/ampere format in the first inference pass + # we no longer need the row-major weight + del self.state.CB + self.weight.data = self.state.CxB + return out + +def quantize_offline(model, bits: int): + assert (bits == 4), f'bits: {bits} is not supported' + + for i, layer in enumerate(model.model.layers): + layer.self_attn.W_pack = bnb.nn.Linear4bit( + layer.self_attn.W_pack.weight.shape[1], + layer.self_attn.W_pack.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + layer.self_attn.o_proj = bnb.nn.Linear4bit( + layer.self_attn.o_proj.weight.shape[1], + layer.self_attn.o_proj.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + + layer.mlp.gate_proj = bnb.nn.Linear4bit( + layer.mlp.gate_proj.weight.shape[1], + layer.mlp.gate_proj.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + layer.mlp.down_proj = bnb.nn.Linear4bit( + layer.mlp.down_proj.weight.shape[1], + layer.mlp.down_proj.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + layer.mlp.up_proj = bnb.nn.Linear4bit( + layer.mlp.up_proj.weight.shape[1], + layer.mlp.up_proj.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + return model + +def quantize_online(model, bits: int): + def quant(weight, bias=None): + if bits == 8: + linear = Linear8bitLtOnline( + weight, + bias, + has_fp16_weights=False, + threshold=6.0, + ) + if bias is not None: + linear.bias = torch.nn.Parameter(bias) + elif bits == 4: + linear = Linear4bitOnline( + weight, + bias, + quant_type="nf4", #fp4/nf4 + ) + else: + raise ValueError("quantize only support 4/8 bit") + return linear + + for i, layer in enumerate(model.model.layers): + layer.self_attn.W_pack = quant(layer.self_attn.W_pack.weight) + layer.self_attn.o_proj = quant(layer.self_attn.o_proj.weight) + layer.mlp.gate_proj = quant(layer.mlp.gate_proj.weight) + layer.mlp.down_proj = quant(layer.mlp.down_proj.weight) + layer.mlp.up_proj = quant(layer.mlp.up_proj.weight) + return model + +def init_model_weight_int4(config, model, state_dict): + #replace Params4bit.cuda with Params4bitCuda + Params4bit.cuda = Params4bitCuda + + for i in range(config.num_hidden_layers): + weight_data = state_dict[f'model.layers.{i}.self_attn.W_pack.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.self_attn.W_pack.weight.quant_state'] + model.model.layers[i].self_attn.W_pack.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + weight_data = state_dict[f'model.layers.{i}.self_attn.o_proj.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.self_attn.o_proj.weight.quant_state'] + model.model.layers[i].self_attn.o_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + weight_data = state_dict[f'model.layers.{i}.mlp.gate_proj.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.mlp.gate_proj.weight.quant_state'] + model.model.layers[i].mlp.gate_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + weight_data = state_dict[f'model.layers.{i}.mlp.up_proj.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.mlp.up_proj.weight.quant_state'] + model.model.layers[i].mlp.up_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + weight_data = state_dict[f'model.layers.{i}.mlp.down_proj.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.mlp.down_proj.weight.quant_state'] + model.model.layers[i].mlp.down_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + model.model.layers[i].input_layernorm.weight = state_dict[f'model.layers.{i}.input_layernorm.weight'] + model.model.layers[i].post_attention_layernorm.weight = state_dict[f'model.layers.{i}.post_attention_layernorm.weight'] + + model.model.embed_tokens.weight = state_dict['model.embed_tokens.weight'] + model.model.norm.weight = state_dict['model.norm.weight'] + model.lm_head.weight = state_dict['lm_head.weight'] + return model \ No newline at end of file diff --git a/ixformer_sdk/train/speedformer/models/bloom/__init__.py b/ixformer_sdk/train/speedformer/models/bloom/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/models/bloom/configuration_bloom.py b/ixformer_sdk/train/speedformer/models/bloom/configuration_bloom.py new file mode 100644 index 00000000..d02df267 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/bloom/configuration_bloom.py @@ -0,0 +1,242 @@ +# coding=utf-8 +# Copyright 2022 the Big Science Workshop and HuggingFace Inc. team. All rights reserved. +# +# 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. +""" Bloom configuration""" +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, List, Mapping, Optional + +from packaging import version + + +if TYPE_CHECKING: + from ... import PreTrainedTokenizer, TensorType + +from transformers.configuration_utils import PretrainedConfig +from transformers.onnx import OnnxConfigWithPast, PatchingSpec +from transformers.utils import is_torch_available, logging + + +logger = logging.get_logger(__name__) + +BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP = { + "bigscience/bloom": "https://huggingface.co/bigscience/bloom/resolve/main/config.json", + "bigscience/bloom-560m": "https://huggingface.co/bigscience/bloom-560m/blob/main/config.json", + "bigscience/bloom-1b1": "https://huggingface.co/bigscience/bloom-1b1/blob/main/config.json", + "bigscience/bloom-1b7": "https://huggingface.co/bigscience/bloom-1b7/blob/main/config.json", + "bigscience/bloom-3b": "https://huggingface.co/bigscience/bloom-3b/blob/main/config.json", + "bigscience/bloom-7b1": "https://huggingface.co/bigscience/bloom-7b1/blob/main/config.json", +} + + +class BloomConfig(PretrainedConfig): + """ + This is the configuration class to store the configuration of a [`BloomModel`]. It is used to instantiate a Bloom + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to the Bloom architecture + [bigscience/bloom](https://huggingface.co/bigscience/bloom). + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 250880): + Vocabulary size of the Bloom model. Defines the maximum number of different tokens that can be represented + by the `inputs_ids` passed when calling [`BloomModel`]. Check [this + discussion](https://huggingface.co/bigscience/bloom/discussions/120#633d28389addb8530b406c2a) on how the + `vocab_size` has been defined. + hidden_size (`int`, *optional*, defaults to 64): + Dimensionality of the embeddings and hidden states. + n_layer (`int`, *optional*, defaults to 2): + Number of hidden layers in the Transformer encoder. + n_head (`int`, *optional*, defaults to 8): + Number of attention heads for each attention layer in the Transformer encoder. + layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): + The epsilon to use in the layer normalization layers. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + apply_residual_connection_post_layernorm (`bool`, *optional*, defaults to `False`): + If enabled, use the layer norm of the hidden states as the residual in the transformer blocks + hidden_dropout (`float`, *optional*, defaults to 0.1): + Dropout rate of the dropout function on the bias dropout. + attention_dropout (`float`, *optional*, defaults to 0.1): + Dropout rate applied to the attention probs + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). + pretraining_tp (`int`, *optional*, defaults to `1`): + Experimental feature. Tensor parallelism rank used during pretraining with Megatron. Please refer to [this + document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is + necessary to ensure exact reproducibility of the pretraining results. Please refer to [this + issue](https://github.com/pytorch/pytorch/issues/76232). Note also that this is enabled only when + `slow_but_exact=True`. + slow_but_exact (`bool`, *optional*, defaults to `False`): + Experimental feature. Whether to use slow but exact implementation of the attention mechanism. While + merging the TP rank tensors, due to slicing operations the results may be slightly different between the + model trained on Megatron and our model. Please refer to [this + issue](https://github.com/pytorch/pytorch/issues/76232). A solution to obtain more accurate results is to + enable this feature. Enabling this will hurt the computational time of the inference. Will be probably + resolved in the future once the main model has been fine-tuned with TP_rank=1. + + Example: + + ```python + >>> from transformers import BloomConfig, BloomModel + + >>> # Initializing a Bloom configuration + >>> configuration = BloomConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = BloomModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "bloom" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "num_hidden_layers": "n_layer", + "num_attention_heads": "n_head", + } + + def __init__( + self, + vocab_size=250880, + hidden_size=64, + n_layer=2, + n_head=8, + layer_norm_epsilon=1e-5, + initializer_range=0.02, + use_cache=True, + bos_token_id=1, + eos_token_id=2, + apply_residual_connection_post_layernorm=False, + hidden_dropout=0.0, + attention_dropout=0.0, + pretraining_tp=1, # TP rank used when training with megatron + slow_but_exact=False, + **kwargs, + ): + self.vocab_size = vocab_size + # Backward compatibility with n_embed kwarg + n_embed = kwargs.pop("n_embed", None) + self.hidden_size = hidden_size if n_embed is None else n_embed + self.n_layer = n_layer + self.n_head = n_head + self.layer_norm_epsilon = layer_norm_epsilon + self.initializer_range = initializer_range + self.use_cache = use_cache + self.pretraining_tp = pretraining_tp + self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm + self.hidden_dropout = hidden_dropout + self.attention_dropout = attention_dropout + + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.slow_but_exact = slow_but_exact + + super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) + + +class BloomOnnxConfig(OnnxConfigWithPast): + torch_onnx_minimum_version = version.parse("1.12") + + def __init__( + self, + config: PretrainedConfig, + task: str = "default", + patching_specs: List[PatchingSpec] = None, + use_past: bool = False, + ): + super().__init__(config, task=task, patching_specs=patching_specs, use_past=use_past) + if not getattr(self._config, "pad_token_id", None): + # TODO: how to do that better? + self._config.pad_token_id = 0 + + @property + def inputs(self) -> Mapping[str, Mapping[int, str]]: + common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}) + if self.use_past: + # BLOOM stores values on dynamic axis 2. For more details see: https://github.com/huggingface/transformers/pull/18344 + self.fill_with_past_key_values_(common_inputs, direction="inputs", inverted_values_shape=True) + common_inputs["attention_mask"] = {0: "batch", 1: "past_sequence + sequence"} + else: + common_inputs["attention_mask"] = {0: "batch", 1: "sequence"} + + return common_inputs + + @property + def num_layers(self) -> int: + return self._config.n_layer + + @property + def num_attention_heads(self) -> int: + return self._config.n_head + + @property + def atol_for_validation(self) -> float: + return 1e-3 + + def generate_dummy_inputs( + self, + tokenizer: "PreTrainedTokenizer", + batch_size: int = -1, + seq_length: int = -1, + is_pair: bool = False, + framework: Optional["TensorType"] = None, + ) -> Mapping[str, Any]: + common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs( + tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework + ) + + # We need to order the input in the way they appears in the forward() + ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]}) + + # Need to add the past_keys + if self.use_past: + if not is_torch_available(): + raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.") + else: + import torch + + batch, seqlen = common_inputs["input_ids"].shape + # Not using the same length for past_key_values + past_key_values_length = seqlen + 2 + head_dim = self._config.hidden_size // self.num_attention_heads + past_key_shape = ( + batch * self.num_attention_heads, + head_dim, + past_key_values_length, + ) + past_value_shape = ( + batch * self.num_attention_heads, + past_key_values_length, + head_dim, + ) + ordered_inputs["past_key_values"] = [ + (torch.zeros(past_key_shape), torch.zeros(past_value_shape)) for _ in range(self.num_layers) + ] + + ordered_inputs["attention_mask"] = common_inputs["attention_mask"] + if self.use_past: + mask_dtype = ordered_inputs["attention_mask"].dtype + ordered_inputs["attention_mask"] = torch.cat( + [ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1 + ) + + return ordered_inputs + + @property + def default_onnx_opset(self) -> int: + return 13 diff --git a/ixformer_sdk/train/speedformer/models/bloom/modeling_attn_mask_utils.py b/ixformer_sdk/train/speedformer/models/bloom/modeling_attn_mask_utils.py new file mode 100644 index 00000000..67555239 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/bloom/modeling_attn_mask_utils.py @@ -0,0 +1,500 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch + + +@dataclass +class AttentionMaskConverter: + """ + A utility attention mask class that allows one to: + - Create a causal 4d mask + - Create a causal 4d mask with slided window + - Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length, + key_value_length) that can be multiplied with attention scores + + Examples: + + ```python + >>> import torch + >>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter + + >>> converter = AttentionMaskConverter(True) + >>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32) + tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]]) + ``` + + Parameters: + is_causal (`bool`): + Whether the attention mask should be a uni-directional (causal) or bi-directional mask. + + sliding_window (`int`, *optional*): + Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer. + """ + + is_causal: bool + sliding_window: int + + def __init__(self, is_causal: bool, sliding_window: Optional[int] = None): + self.is_causal = is_causal + self.sliding_window = sliding_window + + if self.sliding_window is not None and self.sliding_window <= 0: + raise ValueError( + f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`" + ) + + def to_causal_4d( + self, + batch_size: int, + query_length: int, + key_value_length: int, + dtype: torch.dtype, + device: Union[torch.device, "str"] = "cpu", + ) -> Optional[torch.Tensor]: + """ + Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative + bias to upper right hand triangular matrix (causal mask). + """ + if not self.is_causal: + raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.") + + # If shape is not cached, create a new causal mask and cache it + input_shape = (batch_size, query_length) + past_key_values_length = key_value_length - query_length + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if input_shape[-1] > 1 or self.sliding_window is not None: + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + + return causal_4d_mask + + def to_4d( + self, + attention_mask_2d: torch.Tensor, + query_length: int, + dtype: torch.dtype, + key_value_length: Optional[int] = None, + ) -> torch.Tensor: + """ + Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length, + key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is + causal, a causal mask will be added. + """ + input_shape = (attention_mask_2d.shape[0], query_length) + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal: + if key_value_length is None: + raise ValueError( + "This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask." + ) + + past_key_values_length = key_value_length - query_length + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=attention_mask_2d.device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + elif self.sliding_window is not None: + raise NotImplementedError("Sliding window is currently only implemented for causal masking") + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to( + attention_mask_2d.device + ) + + if causal_4d_mask is not None: + expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min) + + # expanded_attn_mask + causal_4d_mask can cause some overflow + expanded_4d_mask = expanded_attn_mask + + return expanded_4d_mask + + @staticmethod + def _make_causal_mask( + input_ids_shape: torch.Size, + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, + ): + """ + Make causal mask used for bi-directional self-attention. + """ + bsz, tgt_len = input_ids_shape + mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device) + mask_cond = torch.arange(mask.size(-1), device=device) + mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0) + + mask = mask.to(dtype) + + if past_key_values_length > 0: + mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1) + + # add lower triangular sliding window mask if necessary + if sliding_window is not None: + diagonal = past_key_values_length - sliding_window + 1 + + context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal) + mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min) + + return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) + + @staticmethod + def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) + + @staticmethod + def _unmask_unattended( + expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float] + ): + # fmt: off + """ + Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when + using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + Details: https://github.com/pytorch/pytorch/issues/110213 + + `expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len]. + `attention_mask` is [bsz, src_seq_len]. + + The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias. + + For example, if `attention_mask` is + ``` + [[0, 0, 1], + [1, 1, 1], + [0, 1, 1]] + ``` + and `expanded_mask` is (e.g. here left-padding case) + ``` + [[[[0, 0, 0], + [0, 0, 0], + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[0, 0, 0], + [0, 1, 0], + [0, 1, 1]]]] + ``` + then the modified `expanded_mask` will be + ``` + [[[[1, 1, 1], <-- modified + [1, 1, 1], <-- modified + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[1, 1, 1], <-- modified + [0, 1, 0], + [0, 1, 1]]]] + ``` + """ + # fmt: on + + # Get the index of the first non-zero value for every sample in the batch. + # In the above example, indices = [[2], [0], [1]]] + tmp = torch.arange(attention_mask.shape[1], 0, -1) + indices = torch.argmax(attention_mask.cpu() * tmp, 1, keepdim=True) + + # Find the batch indexes that have unattended tokens on the leftmost side (e.g. [0, 0, 1, 1, 1]), for which the first rows of the + # expanded mask will be completely unattended. + left_masked_rows = torch.where(indices > 0)[0] + + if left_masked_rows.shape[0] == 0: + return expanded_mask + indices = indices[left_masked_rows] + + max_len = torch.max(indices) + range_tensor = torch.arange(max_len).unsqueeze(0) + range_tensor = range_tensor.repeat(indices.size(0), 1) + + # Avoid unmasking tokens at relevant target positions (on the row axis), by rather unmasking possibly several times the first row that should always be unmasked as we filtered out the batch above. + range_tensor[range_tensor >= indices] = 0 + + # TODO: we may drop support for 3D attention mask as the refactor from Patrick maybe dropped this case + if expanded_mask.dim() == 4: + num_masks = expanded_mask.shape[1] + if num_masks == 1: + # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len] + mask_slice = (left_masked_rows[:, None], 0, range_tensor) + else: + # Broadcast [left_masked_rows, 1, 1], [1, num_masks, 1], [left_masked_rows, 1, max_len] + mask_slice = ( + left_masked_rows[:, None, None], + torch.arange(num_masks)[None, :, None], + range_tensor[:, None, :], + ) + else: + # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len] + mask_slice = (left_masked_rows[:, None], range_tensor) + + expanded_mask[mask_slice] = unmasked_value + + return expanded_mask + + +def _prepare_4d_causal_attention_mask( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + attention_mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + inputs_embeds (`torch.Tensor`): + The embedded inputs as a torch Tensor. + past_key_values_length (`int`): + The length of the key value cache. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + + # 4d mask is passed through the layers + if attention_mask is not None and len(attention_mask.shape) == 2: + attention_mask = attn_mask_converter.to_4d( + attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype + ) + elif attention_mask is not None and len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + else: + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + + return attention_mask + + +# Adapted from _prepare_4d_causal_attention_mask +def _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """ + Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`. + + In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and + `key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks, + allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed). + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + batch_size, query_length = input_shape + + # torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: Fix this as well when using torchdynamo with fullgraph=True. + is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy) + + if attention_mask is not None: + # 4d mask is passed through + if len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype) + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + return attention_mask + + elif not is_tracing and torch.all(attention_mask == 1): + if query_length == 1: + # For query_length == 1, causal attention and bi-directional attention are the same. + attention_mask = None + elif key_value_length == query_length: + attention_mask = None + else: + # Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation + # may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + pass + elif query_length > 1 and key_value_length != query_length: + # See the comment above (https://github.com/pytorch/pytorch/issues/108108). + # Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`. + attention_mask = True + elif is_tracing: + raise ValueError( + 'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.' + ) + + if attention_mask is None: + expanded_4d_mask = None + elif attention_mask is True: + expanded_4d_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + else: + expanded_4d_mask = attn_mask_converter.to_4d( + attention_mask, + input_shape[-1], + dtype=inputs_embeds.dtype, + key_value_length=key_value_length, + ) + + # From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend + # produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213 + # + # This fix is not applied in case we are tracing with torch.jit.trace or symbolic_trace, as _unmask_unattended has a data-dependent + # controlflow that can not be captured properly. + # TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case. + if query_length > 1 and not is_tracing: + expanded_4d_mask = AttentionMaskConverter._unmask_unattended( + expanded_4d_mask, attention_mask, unmasked_value=0.0 + ) + + return expanded_4d_mask + + +def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + batch_size, key_value_length = mask.shape + tgt_len = tgt_len if tgt_len is not None else key_value_length + + # torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: Fix this as well when using torchdynamo with fullgraph=True. + is_tracing = torch.jit.is_tracing() + + if torch.all(mask == 1): + if is_tracing: + pass + elif tgt_len == 1: + # For query_length == 1, causal attention and bi-directional attention are the same. + return None + elif key_value_length == tgt_len: + return None + else: + # Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation + # may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + else: + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _create_4d_causal_attention_mask( + input_shape: Union[torch.Size, Tuple, List], + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, +) -> Optional[torch.Tensor]: + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` + + Args: + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + device (`int`): + The torch device the created mask shall have. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = past_key_values_length + input_shape[-1] + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device + ) + + return attention_mask diff --git a/ixformer_sdk/train/speedformer/models/bloom/modeling_bloom.py b/ixformer_sdk/train/speedformer/models/bloom/modeling_bloom.py new file mode 100644 index 00000000..12ea48d4 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/bloom/modeling_bloom.py @@ -0,0 +1,1250 @@ +# coding=utf-8 +# Copyright 2022 HuggingFace Inc. team and BigScience workshop. +# +# 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. +"""PyTorch BLOOM model.""" + +import math +import warnings +from typing import Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss +from torch.nn import functional as F + +from transformers.file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward +from .modeling_attn_mask_utils import _prepare_4d_causal_attention_mask +from transformers.modeling_outputs import ( + BaseModelOutputWithPastAndCrossAttentions, + CausalLMOutputWithCrossAttentions, + QuestionAnsweringModelOutput, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from .configuration_bloom import BloomConfig + + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "bigscience/bloom-560m" +_CONFIG_FOR_DOC = "BloomConfig" + +BLOOM_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "bigscience/bigscience-small-testing", + "bigscience/bloom-560m", + "bigscience/bloom-1b1", + "bigscience/bloom-1b7", + "bigscience/bloom-3b", + "bigscience/bloom-7b1", + "bigscience/bloom", +] + + +def build_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor: + """ + Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it + relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value + `softmax(l+a) = softmax(l)`. Based on + https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742 + TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly. + + Args: + Returns tensor shaped (batch_size * num_heads, 1, max_seq_len) + attention_mask (`torch.Tensor`): + Token-wise attention mask, this should be of shape (batch_size, max_seq_len). + num_heads (`int`, *required*): + number of heads + dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`): + dtype of the output tensor + """ + batch_size, seq_length = attention_mask.shape + closest_power_of_2 = 2 ** math.floor(math.log2(num_heads)) + base = torch.tensor( + 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32 + ) + powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32) + slopes = torch.pow(base, powers) + + if closest_power_of_2 != num_heads: + extra_base = torch.tensor( + 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32 + ) + num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2) + extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32) + slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0) + + # Note: alibi will added to the attention bias that will be applied to the query, key product of attention + # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length) + # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length) + # => the query_length dimension will then be broadcasted correctly + # This is more or less identical to T5's relative position bias: + # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527 + arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :] + alibi = slopes[..., None] * arange_tensor + return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype) + + +def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor: + """ + Dropout add function + + Args: + x (`torch.tensor`, *required*): + input tensor + residual (`torch.tensor`, *required*): + residual tensor + prob (`float`, *required*): + dropout probability + training (`bool`, *required*): + training mode + """ + out = F.dropout(x, p=prob, training=training) + out = residual + out + return out + + +def bloom_gelu_forward(x: torch.Tensor) -> torch.Tensor: + """ + Custom bias GELU function. Adapted from Megatron-DeepSpeed code. Here we use a simple implementation (inference) to + make the model jitable. + + Args: + x (`torch.tensor`, *required*): + input hidden states + """ + return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))) + + +def bloom_gelu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """ + gradient of tanh approximation of gelu gradient of actual gelu is: 0.5 * (1. + torch.erf(x * 0.70710678)) + + 0.3989423 * x * torch.exp(-0.5 * x * x) + + Args: + g (`torch.tensor`, *required*): + gradient output tensor + x (`torch.tensor`, *required*): + input tensor + """ + x = x[0] # x is a tuple of 1 element, needs to unpack it first + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out) + return ff * g + + +class GeLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(input) + return bloom_gelu_forward(input) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: + input = ctx.saved_tensors + tmp = bloom_gelu_back(grad_output, input) + return tmp + + +class BloomGelu(nn.Module): + """ + BloomBiasGelu wrapper function that make use of the simple function on inference mode to make the model + torchscriptable and use the autograd function in training mode to get the accurate results of the gradients Partly + copied from Megatron-DeepSpeed code and adapted for our needs + + See here why autograd functions are not torchscriptable: https://github.com/pytorch/pytorch/issues/22329 + """ + + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.training: + return GeLUFunction.apply(x) + else: + return bloom_gelu_forward(x) + + +class BloomAttention(nn.Module): + def __init__(self, config: BloomConfig): + super().__init__() + + self.pretraining_tp = config.pretraining_tp + self.slow_but_exact = config.slow_but_exact + + self.hidden_size = config.hidden_size + self.num_heads = config.n_head + self.head_dim = self.hidden_size // self.num_heads + self.split_size = self.hidden_size + self.hidden_dropout = config.hidden_dropout + + if self.head_dim * self.num_heads != self.hidden_size: + raise ValueError( + f"`hidden_size` must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:" + f" {self.num_heads})." + ) + + # Layer-wise attention scaling + self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim) + self.beta = 1.0 + + self.query_key_value = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=True) + self.dense = nn.Linear(self.hidden_size, self.hidden_size) + self.attention_dropout = nn.Dropout(config.attention_dropout) + + def _split_heads(self, fused_qkv: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Split the last dimension into (num_heads, head_dim) without making any copies, results share same memory + storage as `fused_qkv` + + Args: + fused_qkv (`torch.tensor`, *required*): [batch_size, seq_length, num_heads * 3 * head_dim] + + Returns: + query: [batch_size, seq_length, num_heads, head_dim] key: [batch_size, seq_length, num_heads, head_dim] + value: [batch_size, seq_length, num_heads, head_dim] + """ + batch_size, seq_length, three_times_hidden_size = fused_qkv.shape + fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads, 3, self.head_dim) + return fused_qkv[..., 0, :], fused_qkv[..., 1, :], fused_qkv[..., 2, :] + + def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: + """ + Merge heads together over the last dimension + + Args: + x (`torch.tensor`, *required*): [batch_size * num_heads, seq_length, head_dim] + + Returns: + torch.tensor: [batch_size, seq_length, num_heads * head_dim] + """ + # What we want to achieve is: + # batch_size * num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads * head_dim + batch_size_and_num_heads, seq_length, _ = x.shape + batch_size = batch_size_and_num_heads // self.num_heads + + # First view to decompose the batch size + # batch_size * num_heads, seq_length, head_dim -> batch_size, num_heads, seq_length, head_dim + x = x.view(batch_size, self.num_heads, seq_length, self.head_dim) + + # batch_size, num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads, head_dim + x = x.permute(0, 2, 1, 3) + + # batch_size, seq_length, num_heads, head_dim -> batch_size, seq_length, num_heads * head_dim + return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim) + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + alibi: torch.Tensor, + attention_mask: torch.Tensor, + layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + head_mask: Optional[torch.Tensor] = None, + use_cache: bool = False, + output_attentions: bool = False, + ): + fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size] + + # 3 x [batch_size, seq_length, num_heads, head_dim] + (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv) + + batch_size, q_length, _, _ = query_layer.shape + + query_layer = query_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim) + key_layer = key_layer.permute(0, 2, 3, 1).reshape(batch_size * self.num_heads, self.head_dim, q_length) + value_layer = value_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim) + if layer_past is not None: + past_key, past_value = layer_past + # concatenate along seq_length dimension: + # - key: [batch_size * self.num_heads, head_dim, kv_length] + # - value: [batch_size * self.num_heads, kv_length, head_dim] + key_layer = torch.cat((past_key, key_layer), dim=2) + value_layer = torch.cat((past_value, value_layer), dim=1) + + _, _, kv_length = key_layer.shape + + if use_cache is True: + present = (key_layer, value_layer) + else: + present = None + + # [batch_size * num_heads, q_length, kv_length] + # we use `torch.Tensor.baddbmm` instead of `torch.baddbmm` as the latter isn't supported by TorchScript v1.11 + matmul_result = alibi.baddbmm( + batch1=query_layer, + batch2=key_layer, + beta=self.beta, + alpha=self.inv_norm_factor, + ) + + # change view to [batch_size, num_heads, q_length, kv_length] + attention_scores = matmul_result.view(batch_size, self.num_heads, q_length, kv_length) + + # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length] + input_dtype = attention_scores.dtype + # `float16` has a minimum value of -65504.0, whereas `bfloat16` and `float32` have a minimum value of `-3.4e+38` + if input_dtype == torch.float16: + attention_scores = attention_scores.to(torch.float) + attn_weights = torch.masked_fill(attention_scores, attention_mask, torch.finfo(attention_scores.dtype).min) + attention_probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(input_dtype) + + # [batch_size, num_heads, q_length, kv_length] + attention_probs = self.attention_dropout(attention_probs) + + if head_mask is not None: + attention_probs = attention_probs * head_mask + + # change view [batch_size x num_heads, q_length, kv_length] + attention_probs_reshaped = attention_probs.view(batch_size * self.num_heads, q_length, kv_length) + + # matmul: [batch_size * num_heads, q_length, head_dim] + context_layer = torch.bmm(attention_probs_reshaped, value_layer) + + # change view [batch_size, q_length, num_heads * head_dim] + context_layer = self._merge_heads(context_layer) + + # aggregate results across tp ranks. See here: https://github.com/pytorch/pytorch/issues/76232 + if self.pretraining_tp > 1 and self.slow_but_exact: + slices = self.hidden_size / self.pretraining_tp + output_tensor = torch.zeros_like(context_layer) + for i in range(self.pretraining_tp): + output_tensor = output_tensor + F.linear( + context_layer[:, :, int(i * slices) : int((i + 1) * slices)], + self.dense.weight[:, int(i * slices) : int((i + 1) * slices)], + ) + else: + output_tensor = self.dense(context_layer) + + output_tensor = dropout_add(output_tensor, residual, self.hidden_dropout, self.training) + + outputs = (output_tensor, present) + if output_attentions: + outputs += (attention_probs,) + + return outputs + + +class BloomMLP(nn.Module): + def __init__(self, config: BloomConfig): + super().__init__() + hidden_size = config.hidden_size + + self.pretraining_tp = config.pretraining_tp + self.slow_but_exact = config.slow_but_exact + self.dense_h_to_4h = nn.Linear(hidden_size, 4 * hidden_size) + self.gelu_impl = BloomGelu() + self.dense_4h_to_h = nn.Linear(4 * hidden_size, hidden_size) + self.hidden_dropout = config.hidden_dropout + + def forward(self, hidden_states: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + hidden_states = self.gelu_impl(self.dense_h_to_4h(hidden_states)) + + if self.pretraining_tp > 1 and self.slow_but_exact: + intermediate_output = torch.zeros_like(residual) + slices = self.dense_4h_to_h.weight.shape[-1] / self.pretraining_tp + for i in range(self.pretraining_tp): + intermediate_output = intermediate_output + F.linear( + hidden_states[:, :, int(i * slices) : int((i + 1) * slices)], + self.dense_4h_to_h.weight[:, int(i * slices) : int((i + 1) * slices)], + ) + else: + intermediate_output = self.dense_4h_to_h(hidden_states) + + output = dropout_add(intermediate_output, residual, self.hidden_dropout, self.training) + + return output + + +class BloomBlock(nn.Module): + def __init__(self, config: BloomConfig): + super().__init__() + hidden_size = config.hidden_size + + self.input_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + self.num_heads = config.n_head + self.self_attention = BloomAttention(config) + self.post_attention_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + + self.mlp = BloomMLP(config) + + self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm + self.hidden_dropout = config.hidden_dropout + + def forward( + self, + hidden_states: torch.Tensor, + alibi: torch.Tensor, + attention_mask: torch.Tensor, + layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + head_mask: Optional[torch.Tensor] = None, + use_cache: bool = False, + output_attentions: bool = False, + ): + # hidden_states: [batch_size, seq_length, hidden_size] + + # Layer norm at the beginning of the transformer layer. + layernorm_output = self.input_layernorm(hidden_states) + + # Layer norm post the self attention. + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = hidden_states + + # Self attention. + attn_outputs = self.self_attention( + layernorm_output, + residual, + layer_past=layer_past, + attention_mask=attention_mask, + alibi=alibi, + head_mask=head_mask, + use_cache=use_cache, + output_attentions=output_attentions, + ) + + attention_output = attn_outputs[0] + + outputs = attn_outputs[1:] + + layernorm_output = self.post_attention_layernorm(attention_output) + + # Get residual + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = attention_output + + # MLP. + output = self.mlp(layernorm_output, residual) + + if use_cache: + outputs = (output,) + outputs + else: + outputs = (output,) + outputs[1:] + + return outputs # hidden_states, present, attentions + + +class BloomPreTrainedModel(PreTrainedModel): + config_class = BloomConfig + base_model_prefix = "transformer" + supports_gradient_checkpointing = True + _no_split_modules = ["BloomBlock"] + _skip_keys_device_placement = "past_key_values" + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights(self, module: nn.Module): + """Initialize the weights.""" + if isinstance(module, nn.Linear): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + @staticmethod + def _convert_to_standard_cache( + past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]], batch_size: int + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]: + """ + Standardizes the format of the cache so as to match most implementations, i.e. to tuple(tuple([batch_size, + num_heads, ...])) + """ + batch_size_times_num_heads, head_dim, seq_length = past_key_value[0][0].shape + num_heads = batch_size_times_num_heads // batch_size + # key: [batch_size * num_heads, head_dim, seq_length] -> [batch_size, num_heads, head_dim, seq_length] + # value: [batch_size * num_heads, seq_length, head_dim] -> [batch_size, num_heads, seq_length, head_dim] + return tuple( + ( + layer_past[0].view(batch_size, num_heads, head_dim, seq_length), + layer_past[1].view(batch_size, num_heads, seq_length, head_dim), + ) + for layer_past in past_key_value + ) + + @staticmethod + def _convert_to_bloom_cache( + past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]], + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]: + """ + Converts the cache to the format expected by Bloom, i.e. to tuple(tuple([batch_size * num_heads, ...])) + """ + batch_size, num_heads, head_dim, seq_length = past_key_value[0][0].shape + batch_size_times_num_heads = batch_size * num_heads + # key: [batch_size, num_heads, head_dim, seq_length] -> [batch_size * num_heads, head_dim, seq_length] + # value: [batch_size, num_heads, seq_length, head_dim] -> [batch_size * num_heads, seq_length, head_dim] + return tuple( + ( + layer_past[0].view(batch_size_times_num_heads, head_dim, seq_length), + layer_past[1].view(batch_size_times_num_heads, seq_length, head_dim), + ) + for layer_past in past_key_value + ) + + +BLOOM_START_DOCSTRING = r""" + + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`BloomConfig`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +BLOOM_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): + `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values[0][0].shape[2]` + (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary. + + If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as + `input_ids`. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.n_layers`): + Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see + `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have + their past given to this model should not be passed as `input_ids` as they have already been computed. + + Each element of `past_key_values` is a tuple (past_key, past_value): + - past_key: [batch_size * num_heads, head_dim, kv_length] + - past_value: [batch_size * num_heads, kv_length, head_dim] + attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + + If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see + `past_key_values`). + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The bare Bloom Model transformer outputting raw hidden-states without any specific head on top.", + BLOOM_START_DOCSTRING, +) +class BloomModel(BloomPreTrainedModel): + def __init__(self, config: BloomConfig): + super().__init__(config) + + self.embed_dim = config.hidden_size + self.num_heads = config.n_head + + # Embedding + LN Embedding + self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim) + self.word_embeddings_layernorm = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) + + # Transformer blocks + self.h = nn.ModuleList([BloomBlock(config) for _ in range(config.num_hidden_layers)]) + + # Final Layer Norm + self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) + + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def build_alibi_tensor(self, attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor: + return build_alibi_tensor(attention_mask, num_heads, dtype) + + def get_input_embeddings(self): + return self.word_embeddings + + def set_input_embeddings(self, new_embeddings: torch.Tensor): + self.word_embeddings = new_embeddings + + @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=BaseModelOutputWithPastAndCrossAttentions, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **deprecated_arguments, + ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]: + if deprecated_arguments.pop("position_ids", False) is not False: + # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None` + warnings.warn( + "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore" + " passing `position_ids`.", + FutureWarning, + ) + if len(deprecated_arguments) > 0: + raise ValueError(f"Got unexpected arguments: {deprecated_arguments}") + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if past_key_values is None: + past_key_values = tuple([None] * len(self.h)) + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape batch_size x num_heads x N x N + # head_mask has shape n_layer x batch x num_heads x N x N + head_mask = self.get_head_mask(head_mask, self.config.n_layer) + + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + + hidden_states = self.word_embeddings_layernorm(inputs_embeds) + + presents = () if use_cache else None + all_self_attentions = () if output_attentions else None + all_hidden_states = () if output_hidden_states else None + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # Compute alibi tensor: check build_alibi_tensor documentation + seq_length_with_past = seq_length + past_key_values_length = 0 + if past_key_values[0] is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + if attention_mask is None: + attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device) + else: + attention_mask = attention_mask.to(hidden_states.device) + + alibi = self.build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype) + + causal_mask = _prepare_4d_causal_attention_mask( + attention_mask, + input_shape=(batch_size, seq_length), + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + causal_mask = causal_mask.bool() + + for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if self.gradient_checkpointing and self.training: + outputs = self._gradient_checkpointing_func( + block.__call__, + hidden_states, + alibi, + causal_mask, + layer_past, + head_mask[i], + use_cache, + output_attentions, + ) + else: + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=causal_mask, + head_mask=head_mask[i], + use_cache=use_cache, + output_attentions=output_attentions, + alibi=alibi, + ) + + hidden_states = outputs[0] + if use_cache is True: + presents = presents + (outputs[1],) + + if output_attentions: + all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],) + + # Add last hidden state + hidden_states = self.ln_f(hidden_states) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None) + + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=presents, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +@add_start_docstrings( + """ + The Bloom Model transformer with a language modeling head on top (linear layer with weights tied to the input + embeddings). + """, + BLOOM_START_DOCSTRING, +) +class BloomForCausalLM(BloomPreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config: BloomConfig): + super().__init__(config) + self.transformer = BloomModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings: torch.Tensor): + self.lm_head = new_embeddings + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + past_key_values: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + **kwargs, + ) -> dict: + # only last tokens for input_ids if past is not None + if past_key_values is not None: + past_length = past_key_values[0][0].shape[2] + + # Some generation methods already pass only the last input ID + if input_ids.shape[1] > past_length: + remove_prefix_length = past_length + else: + # Default to old behavior: keep only final ID + remove_prefix_length = input_ids.shape[1] - 1 + + input_ids = input_ids[:, remove_prefix_length:] + + # the cache may be in the stardard format (e.g. in contrastive search), convert to bloom's format if needed + if past_key_values[0][0].shape[0] == input_ids.shape[0]: + past_key_values = self._convert_to_bloom_cache(past_key_values) + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + } + ) + return model_inputs + + @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=CausalLMOutputWithCrossAttentions, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **deprecated_arguments, + ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set + `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` + are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` + """ + if deprecated_arguments.pop("position_ids", False) is not False: + # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None` + warnings.warn( + "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore" + " passing `position_ids`.", + FutureWarning, + ) + if len(deprecated_arguments) > 0: + raise ValueError(f"Got unexpected arguments: {deprecated_arguments}") + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + + lm_logits = self.lm_head(hidden_states) + + loss = None + if labels is not None: + # move labels to correct device to enable model parallelism + labels = labels.to(lm_logits.device) + # Shift so that tokens < n predict n + shift_logits = lm_logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + batch_size, seq_length, vocab_size = shift_logits.shape + # Flatten the tokens + loss_fct = CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length) + ) + + if not return_dict: + output = (lm_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutputWithCrossAttentions( + loss=loss, + logits=lm_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + def _reorder_cache( + self, past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: + """ + This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or + [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct + beam_idx at every generation step. + + Output shares the same memory storage as `past`. + """ + standardized_past = self._convert_to_standard_cache(past, batch_size=len(beam_idx)) + + # Get a copy of `beam_idx` on all the devices where we need those indices. + device_to_beam_idx = { + past_state.device: beam_idx.to(past_state.device) for layer_past in past for past_state in layer_past + } + reordered_past = tuple( + ( + layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]), + layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]), + ) + for layer_past in standardized_past + ) + return self._convert_to_bloom_cache(reordered_past) + + +@add_start_docstrings( + """ + The Bloom Model transformer with a sequence classification head on top (linear layer). + + [`BloomForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-1) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + BLOOM_START_DOCSTRING, +) +class BloomForSequenceClassification(BloomPreTrainedModel): + def __init__(self, config: BloomConfig): + super().__init__(config) + self.num_labels = config.num_labels + self.transformer = BloomModel(config) + self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=SequenceClassifierOutputWithPast, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **deprecated_arguments, + ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + if deprecated_arguments.pop("position_ids", False) is not False: + # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None` + warnings.warn( + "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore" + " passing `position_ids`.", + FutureWarning, + ) + if len(deprecated_arguments) > 0: + raise ValueError(f"Got unexpected arguments: {deprecated_arguments}") + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + logger.warning( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@add_start_docstrings( + """ + Bloom Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for + Named-Entity-Recognition (NER) tasks. + """, + BLOOM_START_DOCSTRING, +) +class BloomForTokenClassification(BloomPreTrainedModel): + def __init__(self, config: BloomConfig): + super().__init__(config) + self.num_labels = config.num_labels + + self.transformer = BloomModel(config) + if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None: + classifier_dropout = config.classifier_dropout + elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=TokenClassifierOutput, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **deprecated_arguments, + ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + if deprecated_arguments.pop("position_ids", False) is not False: + # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None` + warnings.warn( + "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore" + " passing `position_ids`.", + FutureWarning, + ) + if len(deprecated_arguments) > 0: + raise ValueError(f"Got unexpected arguments: {deprecated_arguments}") + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + hidden_states = self.dropout(hidden_states) + logits = self.classifier(hidden_states) + + loss = None + if labels is not None: + # move labels to correct device to enable model parallelism + labels = labels.to(logits.device) + batch_size, seq_length = labels.shape + loss_fct = CrossEntropyLoss() + loss = loss_fct( + logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length) + ) + + if not return_dict: + output = (logits,) + transformer_outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@add_start_docstrings( + """ + The BLOOM Model transformer with a span classification head on top for extractive question-answering tasks like + SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`). + """, + BLOOM_START_DOCSTRING, +) +class BloomForQuestionAnswering(BloomPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.transformer = BloomModel(config) + self.qa_outputs = nn.Linear(config.hidden_size, 2) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING.format("batch_size, sequence_length")) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + end_positions: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, QuestionAnsweringModelOutput]: + r""" + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return QuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/ixformer_sdk/train/speedformer/models/chatglm/__init__.py b/ixformer_sdk/train/speedformer/models/chatglm/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/models/chatglm/configuration_chatglm.py b/ixformer_sdk/train/speedformer/models/chatglm/configuration_chatglm.py new file mode 100644 index 00000000..ec32e66d --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/chatglm/configuration_chatglm.py @@ -0,0 +1,61 @@ +from transformers import PretrainedConfig + + +class ChatGLMConfig(PretrainedConfig): + model_type = "chatglm" + def __init__( + self, + num_layers=28, + padded_vocab_size=65024, + hidden_size=4096, + ffn_hidden_size=13696, + kv_channels=128, + num_attention_heads=32, + seq_length=2048, + hidden_dropout=0.0, + classifier_dropout=None, + attention_dropout=0.0, + layernorm_epsilon=1e-5, + rmsnorm=True, + apply_residual_connection_post_layernorm=False, + post_layer_norm=True, + add_bias_linear=False, + add_qkv_bias=False, + bias_dropout_fusion=True, + multi_query_attention=False, + multi_query_group_num=1, + apply_query_key_layer_scaling=True, + attention_softmax_in_fp32=True, + fp32_residual_connection=False, + quantization_bit=0, + pre_seq_len=None, + prefix_projection=False, + **kwargs + ): + self.num_layers = num_layers + self.vocab_size = padded_vocab_size + self.padded_vocab_size = padded_vocab_size + self.hidden_size = hidden_size + self.ffn_hidden_size = ffn_hidden_size + self.kv_channels = kv_channels + self.num_attention_heads = num_attention_heads + self.seq_length = seq_length + self.hidden_dropout = hidden_dropout + self.classifier_dropout = classifier_dropout + self.attention_dropout = attention_dropout + self.layernorm_epsilon = layernorm_epsilon + self.rmsnorm = rmsnorm + self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm + self.post_layer_norm = post_layer_norm + self.add_bias_linear = add_bias_linear + self.add_qkv_bias = add_qkv_bias + self.bias_dropout_fusion = bias_dropout_fusion + self.multi_query_attention = multi_query_attention + self.multi_query_group_num = multi_query_group_num + self.apply_query_key_layer_scaling = apply_query_key_layer_scaling + self.attention_softmax_in_fp32 = attention_softmax_in_fp32 + self.fp32_residual_connection = fp32_residual_connection + self.quantization_bit = quantization_bit + self.pre_seq_len = pre_seq_len + self.prefix_projection = prefix_projection + super().__init__(**kwargs) diff --git a/ixformer_sdk/train/speedformer/models/chatglm/modeling_chatglm.py b/ixformer_sdk/train/speedformer/models/chatglm/modeling_chatglm.py new file mode 100644 index 00000000..4f987a35 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/chatglm/modeling_chatglm.py @@ -0,0 +1,1300 @@ +""" PyTorch ChatGLM model. """ + +import math +import copy +import warnings +import re +import sys + +import torch +import torch.utils.checkpoint +import torch.nn.functional as F +from torch import nn +from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss +from torch.nn.utils import skip_init +from typing import Optional, Tuple, Union, List, Callable, Dict, Any +from copy import deepcopy + +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.generation.logits_process import LogitsProcessor +from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput + +from .configuration_chatglm import ChatGLMConfig + +# flags required to enable jit fusion kernels + +if sys.platform != 'darwin': + torch._C._jit_set_profiling_mode(False) + torch._C._jit_set_profiling_executor(False) + torch._C._jit_override_can_fuse_on_cpu(True) + torch._C._jit_override_can_fuse_on_gpu(True) + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "THUDM/ChatGLM" +_CONFIG_FOR_DOC = "ChatGLMConfig" + +CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "THUDM/chatglm3-6b", + # See all ChatGLM models at https://huggingface.co/models?filter=chatglm +] + + +def default_init(cls, *args, **kwargs): + return cls(*args, **kwargs) + + +class InvalidScoreLogitsProcessor(LogitsProcessor): + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + if torch.isnan(scores).any() or torch.isinf(scores).any(): + scores.zero_() + scores[..., 5] = 5e4 + return scores + + +class PrefixEncoder(torch.nn.Module): + """ + The torch.nn model to encode the prefix + Input shape: (batch-size, prefix-length) + Output shape: (batch-size, prefix-length, 2*layers*hidden) + """ + + def __init__(self, config: ChatGLMConfig): + super().__init__() + self.prefix_projection = config.prefix_projection + if self.prefix_projection: + # Use a two-layer MLP to encode the prefix + kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2 + self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size) + self.trans = torch.nn.Sequential( + torch.nn.Linear(kv_size, config.hidden_size), + torch.nn.Tanh(), + torch.nn.Linear(config.hidden_size, kv_size) + ) + else: + self.embedding = torch.nn.Embedding(config.pre_seq_len, + config.num_layers * config.kv_channels * config.multi_query_group_num * 2) + + def forward(self, prefix: torch.Tensor): + if self.prefix_projection: + prefix_tokens = self.embedding(prefix) + past_key_values = self.trans(prefix_tokens) + else: + past_key_values = self.embedding(prefix) + return past_key_values + + +def split_tensor_along_last_dim( + tensor: torch.Tensor, + num_partitions: int, + contiguous_split_chunks: bool = False, +) -> List[torch.Tensor]: + """Split a tensor along its last dimension. + + Arguments: + tensor: input tensor. + num_partitions: number of partitions to split the tensor + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + + Returns: + A list of Tensors + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + last_dim_size = tensor.size()[last_dim] // num_partitions + # Split. + tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) + # Note: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tensor_list + + +class RotaryEmbedding(nn.Module): + def __init__(self, dim, original_impl=False, device=None, dtype=None): + super().__init__() + inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim)) + self.register_buffer("inv_freq", inv_freq) + self.dim = dim + self.original_impl = original_impl + + def forward_impl( + self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000 + ): + """Enhanced Transformer with Rotary Position Embedding. + + Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/ + transformers/rope/__init__.py. MIT License: + https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license. + """ + # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$ + theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=torch.float, device=device) / n_elem)) + + # Create position indexes `[0, 1, ..., seq_len - 1]` + seq_idx = torch.arange(seq_len, dtype=torch.float, device=device) + + # Calculate the product of position index and $\theta_i$ + idx_theta = torch.outer(seq_idx, theta).float() + + cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1) + + # this is to mimic the behaviour of complex32, else we will get different results + if dtype in (torch.float16, torch.bfloat16, torch.int8): + cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half() + return cache + + def forward(self, max_seq_len, offset=0): + return self.forward_impl( + max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device + ) + + +@torch.jit.script +def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor: + # x: [sq, b, np, hn] + sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3) + rot_dim = rope_cache.shape[-2] * 2 + x, x_pass = x[..., :rot_dim], x[..., rot_dim:] + # truncate to support variable sizes + rope_cache = rope_cache[:sq] + xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2) + rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2) + x_out2 = torch.stack( + [ + xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1], + xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1], + ], + -1, + ) + x_out2 = x_out2.flatten(3) + return torch.cat((x_out2, x_pass), dim=-1) + + +class RMSNorm(torch.nn.Module): + def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs): + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor): + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + + return (self.weight * hidden_states).to(input_dtype) + + +class CoreAttention(torch.nn.Module): + def __init__(self, config: ChatGLMConfig, layer_number): + super(CoreAttention, self).__init__() + + self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling + self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32 + if self.apply_query_key_layer_scaling: + self.attention_softmax_in_fp32 = True + self.layer_number = max(1, layer_number) + + projection_size = config.kv_channels * config.num_attention_heads + + # Per attention head and per partition values. + self.hidden_size_per_partition = projection_size + self.hidden_size_per_attention_head = projection_size // config.num_attention_heads + self.num_attention_heads_per_partition = config.num_attention_heads + + coeff = None + self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) + if self.apply_query_key_layer_scaling: + coeff = self.layer_number + self.norm_factor *= coeff + self.coeff = coeff + + self.attention_dropout = torch.nn.Dropout(config.attention_dropout) + + def forward(self, query_layer, key_layer, value_layer, attention_mask): + pytorch_major_version = int(torch.__version__.split('.')[0]) + if pytorch_major_version >= 2: + query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]] + if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]: + context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, + is_causal=True) + else: + if attention_mask is not None: + attention_mask = ~attention_mask + context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, + attention_mask) + context_layer = context_layer.permute(2, 0, 1, 3) + new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) + context_layer = context_layer.reshape(*new_context_layer_shape) + else: + # Raw attention scores + + # [b, np, sq, sk] + output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0)) + + # [sq, b, np, hn] -> [sq, b * np, hn] + query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1) + # [sk, b, np, hn] -> [sk, b * np, hn] + key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1) + + # preallocting input tensor: [b * np, sq, sk] + matmul_input_buffer = torch.empty( + output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype, + device=query_layer.device + ) + + # Raw attention scores. [b * np, sq, sk] + matmul_result = torch.baddbmm( + matmul_input_buffer, + query_layer.transpose(0, 1), # [b * np, sq, hn] + key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] + beta=0.0, + alpha=(1.0 / self.norm_factor), + ) + + # change view to [b, np, sq, sk] + attention_scores = matmul_result.view(*output_size) + + # =========================== + # Attention probs and dropout + # =========================== + + # attention scores and attention mask [b, np, sq, sk] + if self.attention_softmax_in_fp32: + attention_scores = attention_scores.float() + if self.coeff is not None: + attention_scores = attention_scores * self.coeff + if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]: + attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3], + device=attention_scores.device, dtype=torch.bool) + attention_mask.tril_() + attention_mask = ~attention_mask + if attention_mask is not None: + attention_scores = attention_scores.masked_fill(attention_mask, float("-inf")) + attention_probs = F.softmax(attention_scores, dim=-1) + attention_probs = attention_probs.type_as(value_layer) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.attention_dropout(attention_probs) + # ========================= + # Context layer. [sq, b, hp] + # ========================= + + # value_layer -> context layer. + # [sk, b, np, hn] --> [b, np, sq, hn] + + # context layer shape: [b, np, sq, hn] + output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3)) + # change view [sk, b * np, hn] + value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1) + # change view [b * np, sq, sk] + attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1) + # matmul: [b * np, sq, hn] + context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1)) + # change view [b, np, sq, hn] + context_layer = context_layer.view(*output_size) + # [b, np, sq, hn] --> [sq, b, np, hn] + context_layer = context_layer.permute(2, 0, 1, 3).contiguous() + # [sq, b, np, hn] --> [sq, b, hp] + new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) + context_layer = context_layer.view(*new_context_layer_shape) + + return context_layer + + +class SelfAttention(torch.nn.Module): + """Parallel self-attention layer abstract class. + + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__(self, config: ChatGLMConfig, layer_number, device=None): + super(SelfAttention, self).__init__() + self.layer_number = max(1, layer_number) + + self.projection_size = config.kv_channels * config.num_attention_heads # 128 * 32 + + # Per attention head and per partition values. + self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads # 128 + self.num_attention_heads_per_partition = config.num_attention_heads # 32 + + self.multi_query_attention = config.multi_query_attention + self.qkv_hidden_size = 3 * self.projection_size + if self.multi_query_attention: + self.num_multi_query_groups_per_partition = config.multi_query_group_num + self.qkv_hidden_size = ( + self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num + ) # 4096 + 2 * 128 * 2 + self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size, + bias=config.add_bias_linear or config.add_qkv_bias, + device=device, **_config_to_kwargs(config) + ) + + self.core_attention = CoreAttention(config, self.layer_number) + + # Output. + self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear, + device=device, **_config_to_kwargs(config) + ) + + def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None): + if self.multi_query_attention: + num_attention_heads = self.num_multi_query_groups_per_partition + else: + num_attention_heads = self.num_attention_heads_per_partition + return torch.empty( + inference_max_sequence_len, + batch_size, + num_attention_heads, + self.hidden_size_per_attention_head, + dtype=dtype, + device=device, + ) + + def forward( + self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True + ): + # hidden_states: [sq, b, h] + + # ================================================= + # Pre-allocate memory for key-values for inference. + # ================================================= + # ===================== + # Query, Key, and Value + # ===================== + + # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)] + mixed_x_layer = self.query_key_value(hidden_states) + + if self.multi_query_attention: + (query_layer, key_layer, value_layer) = mixed_x_layer.split( + [ + self.num_attention_heads_per_partition * self.hidden_size_per_attention_head, + self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, + self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, + ], + dim=-1, + ) + query_layer = query_layer.view( + query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + key_layer = key_layer.view( + key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head) + ) + value_layer = value_layer.view( + value_layer.size()[:-1] + + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head) + ) + else: + new_tensor_shape = mixed_x_layer.size()[:-1] + \ + (self.num_attention_heads_per_partition, + 3 * self.hidden_size_per_attention_head) + mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) + + # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn] + (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3) + + # apply relative positional encoding (rotary embedding) + if rotary_pos_emb is not None: + query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb) + key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb) + + # adjust key and value for inference + if kv_cache is not None: + cache_k, cache_v = kv_cache + key_layer = torch.cat((cache_k, key_layer), dim=0) + value_layer = torch.cat((cache_v, value_layer), dim=0) + if use_cache: + kv_cache = (key_layer, value_layer) + else: + kv_cache = None + + if self.multi_query_attention: + key_layer = key_layer.unsqueeze(-2) + key_layer = key_layer.expand( + -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1 + ) + key_layer = key_layer.contiguous().view( + key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + value_layer = value_layer.unsqueeze(-2) + value_layer = value_layer.expand( + -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1 + ) + value_layer = value_layer.contiguous().view( + value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + + # ================================== + # core attention computation + # ================================== + + context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) + + # ================= + # Output. [sq, b, h] + # ================= + + output = self.dense(context_layer) + + return output, kv_cache + + +def _config_to_kwargs(args): + common_kwargs = { + "dtype": args.torch_dtype, + } + return common_kwargs + + +class MLP(torch.nn.Module): + """MLP. + + MLP will take the input with h hidden state, project it to 4*h + hidden dimension, perform nonlinear transformation, and project the + state back into h hidden dimension. + """ + + def __init__(self, config: ChatGLMConfig, device=None): + super(MLP, self).__init__() + + self.add_bias = config.add_bias_linear + + # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf + self.dense_h_to_4h = nn.Linear( + config.hidden_size, + config.ffn_hidden_size * 2, + bias=self.add_bias, + device=device, + **_config_to_kwargs(config) + ) + + def swiglu(x): + x = torch.chunk(x, 2, dim=-1) + return F.silu(x[0]) * x[1] + + self.activation_func = swiglu + + # Project back to h. + self.dense_4h_to_h = nn.Linear( + config.ffn_hidden_size, + config.hidden_size, + bias=self.add_bias, + device=device, + **_config_to_kwargs(config) + ) + + def forward(self, hidden_states): + # [s, b, 4hp] + intermediate_parallel = self.dense_h_to_4h(hidden_states) + intermediate_parallel = self.activation_func(intermediate_parallel) + # [s, b, h] + output = self.dense_4h_to_h(intermediate_parallel) + return output + + +class GLMBlock(torch.nn.Module): + """A single transformer layer. + + Transformer layer takes input with size [s, b, h] and returns an + output of the same size. + """ + + def __init__(self, config: ChatGLMConfig, layer_number, device=None): + super(GLMBlock, self).__init__() + self.layer_number = layer_number + + self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm + + self.fp32_residual_connection = config.fp32_residual_connection + + LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm + # Layernorm on the input data. + self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, + dtype=config.torch_dtype) + + # Self attention. + self.self_attention = SelfAttention(config, layer_number, device=device) + self.hidden_dropout = config.hidden_dropout + + # Layernorm on the attention output + self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, + dtype=config.torch_dtype) + + # MLP + self.mlp = MLP(config, device=device) + + def forward( + self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True, + ): + # hidden_states: [s, b, h] + + # Layer norm at the beginning of the transformer layer. + layernorm_output = self.input_layernorm(hidden_states) + # Self attention. + attention_output, kv_cache = self.self_attention( + layernorm_output, + attention_mask, + rotary_pos_emb, + kv_cache=kv_cache, + use_cache=use_cache + ) + + # Residual connection. + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = hidden_states + + layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training) + layernorm_input = residual + layernorm_input + + # Layer norm post the self attention. + layernorm_output = self.post_attention_layernorm(layernorm_input) + + # MLP. + mlp_output = self.mlp(layernorm_output) + + # Second residual connection. + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = layernorm_input + + output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training) + output = residual + output + + return output, kv_cache + + +class GLMTransformer(torch.nn.Module): + """Transformer class.""" + + def __init__(self, config: ChatGLMConfig, device=None): + super(GLMTransformer, self).__init__() + + self.fp32_residual_connection = config.fp32_residual_connection + self.post_layer_norm = config.post_layer_norm + + # Number of layers. + self.num_layers = config.num_layers + + # Transformer layers. + def build_layer(layer_number): + return GLMBlock(config, layer_number, device=device) + + self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)]) + + if self.post_layer_norm: + LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm + # Final layer norm before output. + self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, + dtype=config.torch_dtype) + + self.gradient_checkpointing = False + + def _get_layer(self, layer_number): + return self.layers[layer_number] + + def forward( + self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None, + use_cache: Optional[bool] = True, + output_hidden_states: Optional[bool] = False, + ): + if not kv_caches: + kv_caches = [None for _ in range(self.num_layers)] + presents = () if use_cache else None + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + all_self_attentions = None + all_hidden_states = () if output_hidden_states else None + for index in range(self.num_layers): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer = self._get_layer(index) + if self.gradient_checkpointing and self.training: + layer_ret = torch.utils.checkpoint.checkpoint( + layer, + hidden_states, + attention_mask, + rotary_pos_emb, + kv_caches[index], + use_cache, + use_reentrant=False + ) + else: + layer_ret = layer( + hidden_states, + attention_mask, + rotary_pos_emb, + kv_cache=kv_caches[index], + use_cache=use_cache + ) + hidden_states, kv_cache = layer_ret + if use_cache: + presents = presents + (kv_cache,) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # Final layer norm. + if self.post_layer_norm: + hidden_states = self.final_layernorm(hidden_states) + + return hidden_states, presents, all_hidden_states, all_self_attentions + + +class ChatGLMPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and + a simple interface for downloading and loading pretrained models. + """ + + is_parallelizable = False + supports_gradient_checkpointing = True + config_class = ChatGLMConfig + base_model_prefix = "transformer" + _no_split_modules = ["GLMBlock"] + + def _init_weights(self, module: nn.Module): + """Initialize the weights.""" + return + + def get_masks(self, input_ids, past_key_values, padding_mask=None): + batch_size, seq_length = input_ids.shape + full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device) + full_attention_mask.tril_() + past_length = 0 + if past_key_values: + past_length = past_key_values[0][0].shape[0] + if past_length: + full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length, + device=input_ids.device), full_attention_mask), dim=-1) + if padding_mask is not None: + full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1) + if not past_length and padding_mask is not None: + full_attention_mask -= padding_mask.unsqueeze(-1) - 1 + full_attention_mask = (full_attention_mask < 0.5).bool() + full_attention_mask.unsqueeze_(1) + return full_attention_mask + + def get_position_ids(self, input_ids, device): + batch_size, seq_length = input_ids.shape + position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1) + return position_ids + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + if not self.supports_gradient_checkpointing: + raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.") + + +class Embedding(torch.nn.Module): + """Language model embeddings.""" + + def __init__(self, config: ChatGLMConfig, device=None): + super(Embedding, self).__init__() + + self.hidden_size = config.hidden_size + # Word embeddings (parallel). + self.word_embeddings = nn.Embedding( + config.padded_vocab_size, + self.hidden_size, + dtype=config.torch_dtype, + device=device + ) + self.fp32_residual_connection = config.fp32_residual_connection + + def forward(self, input_ids): + # Embeddings. + words_embeddings = self.word_embeddings(input_ids) + embeddings = words_embeddings + # Data format change to avoid explicit tranposes : [b s h] --> [s b h]. + embeddings = embeddings.transpose(0, 1).contiguous() + # If the input flag for fp32 residual connection is set, convert for float. + if self.fp32_residual_connection: + embeddings = embeddings.float() + return embeddings + + +class ChatGLMModel(ChatGLMPreTrainedModel): + def __init__(self, config: ChatGLMConfig, device=None, empty_init=True): + super().__init__(config) + if empty_init: + init_method = skip_init + else: + init_method = default_init + init_kwargs = {} + if device is not None: + init_kwargs["device"] = device + self.embedding = init_method(Embedding, config, **init_kwargs) + self.num_layers = config.num_layers + self.multi_query_group_num = config.multi_query_group_num + self.kv_channels = config.kv_channels + + # Rotary positional embeddings + self.seq_length = config.seq_length + rotary_dim = ( + config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels + ) + + self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device, + dtype=config.torch_dtype) + self.encoder = init_method(GLMTransformer, config, **init_kwargs) + self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False, + dtype=config.torch_dtype, **init_kwargs) + self.pre_seq_len = config.pre_seq_len + self.prefix_projection = config.prefix_projection + if self.pre_seq_len is not None: + for param in self.parameters(): + param.requires_grad = False + self.prefix_tokens = torch.arange(self.pre_seq_len).long() + self.prefix_encoder = PrefixEncoder(config) + self.dropout = torch.nn.Dropout(0.1) + + def get_input_embeddings(self): + return self.embedding.word_embeddings + + def set_input_embeddings(self, value): + self.embedding.word_embeddings = value + + def get_prompt(self, batch_size, device, dtype=torch.half): + prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device) + past_key_values = self.prefix_encoder(prefix_tokens).type(dtype) + past_key_values = past_key_values.view( + batch_size, + self.pre_seq_len, + self.num_layers * 2, + self.multi_query_group_num, + self.kv_channels + ) + # seq_len, b, nh, hidden_size + past_key_values = self.dropout(past_key_values) + past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2) + return past_key_values + + def forward( + self, + input_ids, + position_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.BoolTensor] = None, + full_attention_mask: Optional[torch.BoolTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + inputs_embeds: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + batch_size, seq_length = input_ids.shape + + if inputs_embeds is None: + inputs_embeds = self.embedding(input_ids) + + if self.pre_seq_len is not None: + if past_key_values is None: + past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device, + dtype=inputs_embeds.dtype) + if attention_mask is not None: + attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)), + attention_mask], dim=-1) + + if full_attention_mask is None: + if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1): + full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask) + + # Rotary positional embeddings + rotary_pos_emb = self.rotary_pos_emb(self.seq_length) + if position_ids is not None: + rotary_pos_emb = rotary_pos_emb[position_ids] + else: + rotary_pos_emb = rotary_pos_emb[None, :seq_length] + rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous() + + # Run encoder. + hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder( + inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb, + kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states + ) + + if not return_dict: + return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=presents, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + def quantize(self, weight_bit_width: int): + from .quantization import quantize + quantize(self.encoder, weight_bit_width) + return self + + +class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel): + def __init__(self, config: ChatGLMConfig, empty_init=True, device=None): + super().__init__(config) + + self.max_sequence_length = config.max_length + self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device) + self.config = config + self.quantized = False + + if self.config.quantization_bit: + self.quantize(self.config.quantization_bit, empty_init=True) + + def _update_model_kwargs_for_generation( + self, + outputs: ModelOutput, + model_kwargs: Dict[str, Any], + is_encoder_decoder: bool = False, + standardize_cache_format: bool = False, + ) -> Dict[str, Any]: + # update past_key_values + model_kwargs["past_key_values"] = self._extract_past_from_model_output( + outputs, standardize_cache_format=standardize_cache_format + ) + + # update attention mask + if "attention_mask" in model_kwargs: + attention_mask = model_kwargs["attention_mask"] + model_kwargs["attention_mask"] = torch.cat( + [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1 + ) + + # update position ids + if "position_ids" in model_kwargs: + position_ids = model_kwargs["position_ids"] + new_position_id = position_ids[..., -1:].clone() + new_position_id += 1 + model_kwargs["position_ids"] = torch.cat( + [position_ids, new_position_id], dim=-1 + ) + + model_kwargs["is_first_forward"] = False + return model_kwargs + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + past_key_values: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + is_first_forward: bool = True, + **kwargs + ) -> dict: + # only last token for input_ids if past is not None + if position_ids is None: + position_ids = self.get_position_ids(input_ids, device=input_ids.device) + if not is_first_forward: + if past_key_values is not None: + position_ids = position_ids[..., -1:] + input_ids = input_ids[:, -1:] + return { + "input_ids": input_ids, + "past_key_values": past_key_values, + "position_ids": position_ids, + "attention_mask": attention_mask, + "return_last_logit": True, + "use_cache": use_cache + } + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Tuple[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + return_last_logit: Optional[bool] = False, + ): + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + if return_last_logit: + hidden_states = hidden_states[-1:] + lm_logits = self.transformer.output_layer(hidden_states) + lm_logits = lm_logits.transpose(0, 1).contiguous() + + loss = None + if labels is not None: + lm_logits = lm_logits.to(torch.float32) + + # Shift so that tokens < n predict n + shift_logits = lm_logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss(ignore_index=-100) + loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + + lm_logits = lm_logits.to(hidden_states.dtype) + loss = loss.to(hidden_states.dtype) + + if not return_dict: + output = (lm_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=lm_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + @staticmethod + def _reorder_cache( + past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: + """ + This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or + [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct + beam_idx at every generation step. + + Output shares the same memory storage as `past`. + """ + return tuple( + ( + layer_past[0].index_select(1, beam_idx.to(layer_past[0].device)), + layer_past[1].index_select(1, beam_idx.to(layer_past[1].device)), + ) + for layer_past in past + ) + + def process_response(self, output, history): + content = "" + history = deepcopy(history) + for response in output.split("<|assistant|>"): + if "\n" in response: + metadata, content = response.split("\n", maxsplit=1) + else: + metadata, content = "", response + if not metadata.strip(): + content = content.strip() + history.append({"role": "assistant", "metadata": metadata, "content": content}) + content = content.replace("[[训练时间]]", "2023年") + else: + history.append({"role": "assistant", "metadata": metadata, "content": content}) + if history[0]["role"] == "system" and "tools" in history[0]: + content = "\n".join(content.split("\n")[1:-1]) + def tool_call(**kwargs): + return kwargs + parameters = eval(content) + content = {"name": metadata.strip(), "parameters": parameters} + else: + content = {"name": metadata.strip(), "content": content} + return content, history + + @torch.inference_mode() + def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user", + max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None, + **kwargs): + if history is None: + history = [] + if logits_processor is None: + logits_processor = LogitsProcessorList() + logits_processor.append(InvalidScoreLogitsProcessor()) + gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p, + "temperature": temperature, "logits_processor": logits_processor, **kwargs} + inputs = tokenizer.build_chat_input(query, history=history, role=role) + inputs = inputs.to(self.device) + eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"), + tokenizer.get_command("<|observation|>")] + outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id) + outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1] + response = tokenizer.decode(outputs) + history.append({"role": role, "content": query}) + response, history = self.process_response(response, history) + return response, history + + @torch.inference_mode() + def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user", + past_key_values=None,max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8, + logits_processor=None, return_past_key_values=False, **kwargs): + if history is None: + history = [] + if logits_processor is None: + logits_processor = LogitsProcessorList() + logits_processor.append(InvalidScoreLogitsProcessor()) + eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"), + tokenizer.get_command("<|observation|>")] + gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p, + "temperature": temperature, "logits_processor": logits_processor, **kwargs} + if past_key_values is None: + inputs = tokenizer.build_chat_input(query, history=history, role=role) + else: + inputs = tokenizer.build_chat_input(query, role=role) + inputs = inputs.to(self.device) + if past_key_values is not None: + past_length = past_key_values[0][0].shape[0] + if self.transformer.pre_seq_len is not None: + past_length -= self.transformer.pre_seq_len + inputs.position_ids += past_length + attention_mask = inputs.attention_mask + attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1) + inputs['attention_mask'] = attention_mask + history.append({"role": role, "content": query}) + for outputs in self.stream_generate(**inputs, past_key_values=past_key_values, + eos_token_id=eos_token_id, return_past_key_values=return_past_key_values, + **gen_kwargs): + if return_past_key_values: + outputs, past_key_values = outputs + outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1] + response = tokenizer.decode(outputs) + if response and response[-1] != "�": + response, new_history = self.process_response(response, history) + if return_past_key_values: + yield response, new_history, past_key_values + else: + yield response, new_history + + @torch.inference_mode() + def stream_generate( + self, + input_ids, + generation_config: Optional[GenerationConfig] = None, + logits_processor: Optional[LogitsProcessorList] = None, + stopping_criteria: Optional[StoppingCriteriaList] = None, + prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None, + return_past_key_values=False, + **kwargs, + ): + batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1] + + if generation_config is None: + generation_config = self.generation_config + generation_config = copy.deepcopy(generation_config) + model_kwargs = generation_config.update(**kwargs) + model_kwargs["use_cache"] = generation_config.use_cache + bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id + + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None + + has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None + if has_default_max_length and generation_config.max_new_tokens is None: + warnings.warn( + f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. " + "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we" + " recommend using `max_new_tokens` to control the maximum length of the generation.", + UserWarning, + ) + elif generation_config.max_new_tokens is not None: + generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length + if not has_default_max_length: + logger.warn( + f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(=" + f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. " + "Please refer to the documentation for more information. " + "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)", + UserWarning, + ) + + if input_ids_seq_length >= generation_config.max_length: + input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" + logger.warning( + f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to" + f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider" + " increasing `max_new_tokens`." + ) + + # 2. Set generation parameters if not already defined + logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() + stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() + + logits_processor = self._get_logits_processor( + generation_config=generation_config, + input_ids_seq_length=input_ids_seq_length, + encoder_input_ids=input_ids, + prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, + logits_processor=logits_processor, + ) + + stopping_criteria = self._get_stopping_criteria( + generation_config=generation_config, stopping_criteria=stopping_criteria + ) + logits_warper = self._get_logits_warper(generation_config) + + unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1) + scores = None + while True: + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + # forward pass to get next token + outputs = self( + **model_inputs, + return_dict=True, + output_attentions=False, + output_hidden_states=False, + ) + + next_token_logits = outputs.logits[:, -1, :] + + # pre-process distribution + next_token_scores = logits_processor(input_ids, next_token_logits) + next_token_scores = logits_warper(input_ids, next_token_scores) + + # sample + probs = nn.functional.softmax(next_token_scores, dim=-1) + if generation_config.do_sample: + next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) + else: + next_tokens = torch.argmax(probs, dim=-1) + # update generated ids, model inputs, and length for next step + input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) + model_kwargs = self._update_model_kwargs_for_generation( + outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder + ) + unfinished_sequences = unfinished_sequences.mul( + next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0) + ) + if return_past_key_values: + yield input_ids, outputs.past_key_values + else: + yield input_ids + # stop when each sentence is finished, or if we exceed the maximum length + if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores): + break + + def quantize(self, bits: int, empty_init=False, device=None, **kwargs): + if bits == 0: + return + + from .quantization import quantize + + if self.quantized: + logger.info("Already quantized.") + return self + + self.quantized = True + + self.config.quantization_bit = bits + + self.transformer.encoder = quantize(self.transformer.encoder, bits, empty_init=empty_init, device=device, + **kwargs) + return self + + +class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel): + def __init__(self, config: ChatGLMConfig, empty_init=True, device=None): + super().__init__(config) + + self.num_labels = config.num_labels + self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device) + + self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=torch.half) + if config.classifier_dropout is not None: + self.dropout = nn.Dropout(config.classifier_dropout) + else: + self.dropout = None + self.config = config + + if self.config.quantization_bit: + self.quantize(self.config.quantization_bit, empty_init=True) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + full_attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + inputs_embeds: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutputWithPast]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + full_attention_mask=full_attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + pooled_hidden_states = hidden_states[-1] + if self.dropout is not None: + pooled_hidden_states = self.dropout(pooled_hidden_states) + logits = self.classifier_head(pooled_hidden_states) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze().float(), labels.squeeze()) + else: + loss = loss_fct(logits.float(), labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels).float(), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits.float(), labels.view(-1, self.num_labels)) + + if not return_dict: + output = (logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) \ No newline at end of file diff --git a/ixformer_sdk/train/speedformer/models/chatglm/modeling_chatglm_flash.py b/ixformer_sdk/train/speedformer/models/chatglm/modeling_chatglm_flash.py new file mode 100644 index 00000000..6b6ec44c --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/chatglm/modeling_chatglm_flash.py @@ -0,0 +1,1386 @@ +""" PyTorch ChatGLM model. """ + +import math +import copy +import warnings +import re +import sys + +import torch +import torch.utils.checkpoint +import torch.nn.functional as F +from torch import nn +from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss +from torch.nn.utils import skip_init +from typing import Optional, Tuple, Union, List, Callable, Dict, Any +from copy import deepcopy + +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging, is_flash_attn_2_available +from transformers.generation.logits_process import LogitsProcessor +from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput + +from .configuration_chatglm import ChatGLMConfig + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + +# flags required to enable jit fusion kernels + +if sys.platform != 'darwin': + torch._C._jit_set_profiling_mode(False) + torch._C._jit_set_profiling_executor(False) + torch._C._jit_override_can_fuse_on_cpu(True) + torch._C._jit_override_can_fuse_on_gpu(True) + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "THUDM/ChatGLM" +_CONFIG_FOR_DOC = "ChatGLMConfig" + +CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "THUDM/chatglm3-6b", + # See all ChatGLM models at https://huggingface.co/models?filter=chatglm +] + + +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def default_init(cls, *args, **kwargs): + return cls(*args, **kwargs) + + +class InvalidScoreLogitsProcessor(LogitsProcessor): + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + if torch.isnan(scores).any() or torch.isinf(scores).any(): + scores.zero_() + scores[..., 5] = 5e4 + return scores + + +class PrefixEncoder(torch.nn.Module): + """ + The torch.nn model to encode the prefix + Input shape: (batch-size, prefix-length) + Output shape: (batch-size, prefix-length, 2*layers*hidden) + """ + + def __init__(self, config: ChatGLMConfig): + super().__init__() + self.prefix_projection = config.prefix_projection + if self.prefix_projection: + # Use a two-layer MLP to encode the prefix + kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2 + self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size) + self.trans = torch.nn.Sequential( + torch.nn.Linear(kv_size, config.hidden_size), + torch.nn.Tanh(), + torch.nn.Linear(config.hidden_size, kv_size) + ) + else: + self.embedding = torch.nn.Embedding(config.pre_seq_len, + config.num_layers * config.kv_channels * config.multi_query_group_num * 2) + + def forward(self, prefix: torch.Tensor): + if self.prefix_projection: + prefix_tokens = self.embedding(prefix) + past_key_values = self.trans(prefix_tokens) + else: + past_key_values = self.embedding(prefix) + return past_key_values + + +def split_tensor_along_last_dim( + tensor: torch.Tensor, + num_partitions: int, + contiguous_split_chunks: bool = False, +) -> List[torch.Tensor]: + """Split a tensor along its last dimension. + + Arguments: + tensor: input tensor. + num_partitions: number of partitions to split the tensor + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + + Returns: + A list of Tensors + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + last_dim_size = tensor.size()[last_dim] // num_partitions + # Split. + tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) + # Note: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tensor_list + + +class RotaryEmbedding(nn.Module): + def __init__(self, dim, original_impl=False, device=None, dtype=None): + super().__init__() + inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim)) + self.register_buffer("inv_freq", inv_freq) + self.dim = dim + self.original_impl = original_impl + + def forward_impl( + self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000 + ): + """Enhanced Transformer with Rotary Position Embedding. + + Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/ + transformers/rope/__init__.py. MIT License: + https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license. + """ + # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$ + theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=torch.float, device=device) / n_elem)) + + # Create position indexes `[0, 1, ..., seq_len - 1]` + seq_idx = torch.arange(seq_len, dtype=torch.float, device=device) + + # Calculate the product of position index and $\theta_i$ + idx_theta = torch.outer(seq_idx, theta).float() + + cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1) + + # this is to mimic the behaviour of complex32, else we will get different results + if dtype in (torch.float16, torch.bfloat16, torch.int8): + cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half() + return cache + + def forward(self, max_seq_len, offset=0): + return self.forward_impl( + max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device + ) + + +@torch.jit.script +def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor: + # x: [sq, b, np, hn] + sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3) + rot_dim = rope_cache.shape[-2] * 2 + x, x_pass = x[..., :rot_dim], x[..., rot_dim:] + # truncate to support variable sizes + rope_cache = rope_cache[:sq] + xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2) + rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2) + x_out2 = torch.stack( + [ + xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1], + xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1], + ], + -1, + ) + x_out2 = x_out2.flatten(3) + return torch.cat((x_out2, x_pass), dim=-1) + + +class RMSNorm(torch.nn.Module): + def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs): + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor): + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + + return (self.weight * hidden_states).to(input_dtype) + + +class CoreAttention(torch.nn.Module): + def __init__(self, config: ChatGLMConfig, layer_number): + super(CoreAttention, self).__init__() + + self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling + self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32 + if self.apply_query_key_layer_scaling: + self.attention_softmax_in_fp32 = True + self.layer_number = max(1, layer_number) + + projection_size = config.kv_channels * config.num_attention_heads + + # Per attention head and per partition values. + self.hidden_size_per_partition = projection_size + self.hidden_size_per_attention_head = projection_size // config.num_attention_heads + self.num_attention_heads_per_partition = config.num_attention_heads + + coeff = None + self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) + if self.apply_query_key_layer_scaling: + coeff = self.layer_number + self.norm_factor *= coeff + self.coeff = coeff + + self.attention_dropout = torch.nn.Dropout(config.attention_dropout) + + + def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = index_first_axis( + key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + + def forward(self, query_layer, key_layer, value_layer, attention_mask): + # q,k,v = [sq, b, np, hn] + pytorch_major_version = int(torch.__version__.split('.')[0]) + if pytorch_major_version >= 2: + if is_flash_attn_2_available: + query_length, batch_size, num_head, head_dim = query_layer.shape + query_layer, key_layer, value_layer = [k.permute(1, 0, 2, 3) for k in [query_layer, key_layer, value_layer]] + if attention_mask is not None: + query_layer, key_layer, value_layer, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( + query_layer, key_layer, value_layer, attention_mask, query_length + ) + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + attn_output_unpad = flash_attn_varlen_func( + query_layer, + key_layer, + value_layer, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=0.0, + softmax_scale=None, + causal=True, + ) + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) + else: + attn_output = flash_attn_func( + query_layer, key_layer, value_layer, 0.0, softmax_scale=None, causal=True + ) + + context_layer = attn_output.permute(1, 0, 2, 3) + + else: + query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]] + if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]: + context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, + is_causal=True) + else: + if attention_mask is not None: + attention_mask = ~attention_mask + + context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, + attention_mask) + context_layer = context_layer.permute(2, 0, 1, 3) + + new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) + context_layer = context_layer.reshape(*new_context_layer_shape) + + else: + # Raw attention scores + + # [b, np, sq, sk] + output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0)) + + # [sq, b, np, hn] -> [sq, b * np, hn] + query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1) + # [sk, b, np, hn] -> [sk, b * np, hn] + key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1) + + # preallocting input tensor: [b * np, sq, sk] + matmul_input_buffer = torch.empty( + output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype, + device=query_layer.device + ) + + # Raw attention scores. [b * np, sq, sk] + matmul_result = torch.baddbmm( + matmul_input_buffer, + query_layer.transpose(0, 1), # [b * np, sq, hn] + key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] + beta=0.0, + alpha=(1.0 / self.norm_factor), + ) + + # change view to [b, np, sq, sk] + attention_scores = matmul_result.view(*output_size) + + # =========================== + # Attention probs and dropout + # =========================== + + # attention scores and attention mask [b, np, sq, sk] + if self.attention_softmax_in_fp32: + attention_scores = attention_scores.float() + if self.coeff is not None: + attention_scores = attention_scores * self.coeff + if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]: + attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3], + device=attention_scores.device, dtype=torch.bool) + attention_mask.tril_() + attention_mask = ~attention_mask + if attention_mask is not None: + attention_scores = attention_scores.masked_fill(attention_mask, float("-inf")) + attention_probs = F.softmax(attention_scores, dim=-1) + attention_probs = attention_probs.type_as(value_layer) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.attention_dropout(attention_probs) + # ========================= + # Context layer. [sq, b, hp] + # ========================= + + # value_layer -> context layer. + # [sk, b, np, hn] --> [b, np, sq, hn] + + # context layer shape: [b, np, sq, hn] + output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3)) + # change view [sk, b * np, hn] + value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1) + # change view [b * np, sq, sk] + attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1) + # matmul: [b * np, sq, hn] + context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1)) + # change view [b, np, sq, hn] + context_layer = context_layer.view(*output_size) + # [b, np, sq, hn] --> [sq, b, np, hn] + context_layer = context_layer.permute(2, 0, 1, 3).contiguous() + # [sq, b, np, hn] --> [sq, b, hp] + new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) + context_layer = context_layer.view(*new_context_layer_shape) + + return context_layer + + +class SelfAttention(torch.nn.Module): + """Parallel self-attention layer abstract class. + + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__(self, config: ChatGLMConfig, layer_number, device=None): + super(SelfAttention, self).__init__() + self.layer_number = max(1, layer_number) + + self.projection_size = config.kv_channels * config.num_attention_heads + + # Per attention head and per partition values. + self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads + self.num_attention_heads_per_partition = config.num_attention_heads + + self.multi_query_attention = config.multi_query_attention + self.qkv_hidden_size = 3 * self.projection_size + if self.multi_query_attention: + self.num_multi_query_groups_per_partition = config.multi_query_group_num + self.qkv_hidden_size = ( + self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num + ) + self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size, + bias=config.add_bias_linear or config.add_qkv_bias, + device=device, **_config_to_kwargs(config) + ) + + self.core_attention = CoreAttention(config, self.layer_number) + + # Output. + self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear, + device=device, **_config_to_kwargs(config) + ) + + def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None): + if self.multi_query_attention: + num_attention_heads = self.num_multi_query_groups_per_partition + else: + num_attention_heads = self.num_attention_heads_per_partition + return torch.empty( + inference_max_sequence_len, + batch_size, + num_attention_heads, + self.hidden_size_per_attention_head, + dtype=dtype, + device=device, + ) + + def forward( + self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True + ): + # hidden_states: [sq, b, h] + + # ================================================= + # Pre-allocate memory for key-values for inference. + # ================================================= + # ===================== + # Query, Key, and Value + # ===================== + + # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)] + mixed_x_layer = self.query_key_value(hidden_states) + + if self.multi_query_attention: + (query_layer, key_layer, value_layer) = mixed_x_layer.split( + [ + self.num_attention_heads_per_partition * self.hidden_size_per_attention_head, + self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, + self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, + ], + dim=-1, + ) + query_layer = query_layer.view( + query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + key_layer = key_layer.view( + key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head) + ) + value_layer = value_layer.view( + value_layer.size()[:-1] + + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head) + ) + else: + new_tensor_shape = mixed_x_layer.size()[:-1] + \ + (self.num_attention_heads_per_partition, + 3 * self.hidden_size_per_attention_head) + mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) + + # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn] + (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3) + + # apply relative positional encoding (rotary embedding) + if rotary_pos_emb is not None: + query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb) + key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb) + + # adjust key and value for inference + if kv_cache is not None: + cache_k, cache_v = kv_cache + key_layer = torch.cat((cache_k, key_layer), dim=0) + value_layer = torch.cat((cache_v, value_layer), dim=0) + if use_cache: + kv_cache = (key_layer, value_layer) + else: + kv_cache = None + + if self.multi_query_attention: + key_layer = key_layer.unsqueeze(-2) + key_layer = key_layer.expand( + -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1 + ) + key_layer = key_layer.contiguous().view( + key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + value_layer = value_layer.unsqueeze(-2) + value_layer = value_layer.expand( + -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1 + ) + value_layer = value_layer.contiguous().view( + value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + + # ================================== + # core attention computation + # ================================== + + context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) + + # ================= + # Output. [sq, b, h] + # ================= + + output = self.dense(context_layer) + + return output, kv_cache + + +def _config_to_kwargs(args): + common_kwargs = { + "dtype": args.torch_dtype, + } + return common_kwargs + + +class MLP(torch.nn.Module): + """MLP. + + MLP will take the input with h hidden state, project it to 4*h + hidden dimension, perform nonlinear transformation, and project the + state back into h hidden dimension. + """ + + def __init__(self, config: ChatGLMConfig, device=None): + super(MLP, self).__init__() + + self.add_bias = config.add_bias_linear + + # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf + self.dense_h_to_4h = nn.Linear( + config.hidden_size, + config.ffn_hidden_size * 2, + bias=self.add_bias, + device=device, + **_config_to_kwargs(config) + ) + + def swiglu(x): + x = torch.chunk(x, 2, dim=-1) + return F.silu(x[0]) * x[1] + + self.activation_func = swiglu + + # Project back to h. + self.dense_4h_to_h = nn.Linear( + config.ffn_hidden_size, + config.hidden_size, + bias=self.add_bias, + device=device, + **_config_to_kwargs(config) + ) + + def forward(self, hidden_states): + # [s, b, 4hp] + intermediate_parallel = self.dense_h_to_4h(hidden_states) + intermediate_parallel = self.activation_func(intermediate_parallel) + # [s, b, h] + output = self.dense_4h_to_h(intermediate_parallel) + return output + + +class GLMBlock(torch.nn.Module): + """A single transformer layer. + + Transformer layer takes input with size [s, b, h] and returns an + output of the same size. + """ + + def __init__(self, config: ChatGLMConfig, layer_number, device=None): + super(GLMBlock, self).__init__() + self.layer_number = layer_number + + self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm + + self.fp32_residual_connection = config.fp32_residual_connection + + LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm + # Layernorm on the input data. + self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, + dtype=config.torch_dtype) + + # Self attention. + self.self_attention = SelfAttention(config, layer_number, device=device) + self.hidden_dropout = config.hidden_dropout + + # Layernorm on the attention output + self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, + dtype=config.torch_dtype) + + # MLP + self.mlp = MLP(config, device=device) + + def forward( + self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True, + ): + # hidden_states: [s, b, h] + + # Layer norm at the beginning of the transformer layer. + layernorm_output = self.input_layernorm(hidden_states) + # Self attention. + attention_output, kv_cache = self.self_attention( + layernorm_output, + attention_mask, + rotary_pos_emb, + kv_cache=kv_cache, + use_cache=use_cache + ) + + # Residual connection. + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = hidden_states + + layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training) + layernorm_input = residual + layernorm_input + + # Layer norm post the self attention. + layernorm_output = self.post_attention_layernorm(layernorm_input) + + # MLP. + mlp_output = self.mlp(layernorm_output) + + # Second residual connection. + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = layernorm_input + + output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training) + output = residual + output + + return output, kv_cache + + +class GLMTransformer(torch.nn.Module): + """Transformer class.""" + + def __init__(self, config: ChatGLMConfig, device=None): + super(GLMTransformer, self).__init__() + + self.fp32_residual_connection = config.fp32_residual_connection + self.post_layer_norm = config.post_layer_norm + + # Number of layers. + self.num_layers = config.num_layers + + # Transformer layers. + def build_layer(layer_number): + return GLMBlock(config, layer_number, device=device) + + self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)]) + + if self.post_layer_norm: + LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm + # Final layer norm before output. + self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, + dtype=config.torch_dtype) + + self.gradient_checkpointing = False + + def _get_layer(self, layer_number): + return self.layers[layer_number] + + def forward( + self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None, + use_cache: Optional[bool] = True, + output_hidden_states: Optional[bool] = False, + ): + if not kv_caches: + kv_caches = [None for _ in range(self.num_layers)] + presents = () if use_cache else None + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + all_self_attentions = None + all_hidden_states = () if output_hidden_states else None + for index in range(self.num_layers): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer = self._get_layer(index) + if self.gradient_checkpointing and self.training: + layer_ret = torch.utils.checkpoint.checkpoint( + layer, + hidden_states, + attention_mask, + rotary_pos_emb, + kv_caches[index], + use_cache + ) + else: + layer_ret = layer( + hidden_states, + attention_mask, + rotary_pos_emb, + kv_cache=kv_caches[index], + use_cache=use_cache + ) + hidden_states, kv_cache = layer_ret + if use_cache: + presents = presents + (kv_cache,) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # Final layer norm. + if self.post_layer_norm: + hidden_states = self.final_layernorm(hidden_states) + + return hidden_states, presents, all_hidden_states, all_self_attentions + + +class ChatGLMPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and + a simple interface for downloading and loading pretrained models. + """ + + is_parallelizable = False + supports_gradient_checkpointing = True + config_class = ChatGLMConfig + base_model_prefix = "transformer" + _no_split_modules = ["GLMBlock"] + + def _init_weights(self, module: nn.Module): + """Initialize the weights.""" + return + + def get_masks(self, input_ids, past_key_values, padding_mask=None): + batch_size, seq_length = input_ids.shape + full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device) + full_attention_mask.tril_() + past_length = 0 + if past_key_values: + past_length = past_key_values[0][0].shape[0] + if past_length: + full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length, + device=input_ids.device), full_attention_mask), dim=-1) + if padding_mask is not None: + full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1) + if not past_length and padding_mask is not None: + full_attention_mask -= padding_mask.unsqueeze(-1) - 1 + full_attention_mask = (full_attention_mask < 0.5).bool() + full_attention_mask.unsqueeze_(1) + return full_attention_mask + + def get_position_ids(self, input_ids, device): + batch_size, seq_length = input_ids.shape + position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1) + return position_ids + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, GLMTransformer): + module.gradient_checkpointing = value + + +class Embedding(torch.nn.Module): + """Language model embeddings.""" + + def __init__(self, config: ChatGLMConfig, device=None): + super(Embedding, self).__init__() + + self.hidden_size = config.hidden_size + # Word embeddings (parallel). + self.word_embeddings = nn.Embedding( + config.padded_vocab_size, + self.hidden_size, + dtype=config.torch_dtype, + device=device + ) + self.fp32_residual_connection = config.fp32_residual_connection + + def forward(self, input_ids): + # Embeddings. + words_embeddings = self.word_embeddings(input_ids) + embeddings = words_embeddings + # Data format change to avoid explicit tranposes : [b s h] --> [s b h]. + embeddings = embeddings.transpose(0, 1).contiguous() + # If the input flag for fp32 residual connection is set, convert for float. + if self.fp32_residual_connection: + embeddings = embeddings.float() + return embeddings + + +class ChatGLMModel(ChatGLMPreTrainedModel): + def __init__(self, config: ChatGLMConfig, device=None, empty_init=True): + super().__init__(config) + if empty_init: + init_method = skip_init + else: + init_method = default_init + init_kwargs = {} + if device is not None: + init_kwargs["device"] = device + self.embedding = init_method(Embedding, config, **init_kwargs) + self.num_layers = config.num_layers + self.multi_query_group_num = config.multi_query_group_num + self.kv_channels = config.kv_channels + + # Rotary positional embeddings + self.seq_length = config.seq_length + rotary_dim = ( + config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels + ) + + self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device, + dtype=config.torch_dtype) + self.encoder = init_method(GLMTransformer, config, **init_kwargs) + self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False, + dtype=config.torch_dtype, **init_kwargs) + self.pre_seq_len = config.pre_seq_len + self.prefix_projection = config.prefix_projection + if self.pre_seq_len is not None: + for param in self.parameters(): + param.requires_grad = False + self.prefix_tokens = torch.arange(self.pre_seq_len).long() + self.prefix_encoder = PrefixEncoder(config) + self.dropout = torch.nn.Dropout(0.1) + + def get_input_embeddings(self): + return self.embedding.word_embeddings + + def get_prompt(self, batch_size, device, dtype=torch.half): + prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device) + past_key_values = self.prefix_encoder(prefix_tokens).type(dtype) + past_key_values = past_key_values.view( + batch_size, + self.pre_seq_len, + self.num_layers * 2, + self.multi_query_group_num, + self.kv_channels + ) + # seq_len, b, nh, hidden_size + past_key_values = self.dropout(past_key_values) + past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2) + return past_key_values + + def forward( + self, + input_ids, + position_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.BoolTensor] = None, + full_attention_mask: Optional[torch.BoolTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + inputs_embeds: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + batch_size, seq_length = input_ids.shape + + if inputs_embeds is None: + inputs_embeds = self.embedding(input_ids) + + if self.pre_seq_len is not None: + if past_key_values is None: + past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device, + dtype=inputs_embeds.dtype) + if attention_mask is not None: + attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)), + attention_mask], dim=-1) + if full_attention_mask is None: + if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1): + full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask) + + # Rotary positional embeddings + rotary_pos_emb = self.rotary_pos_emb(self.seq_length) + if position_ids is not None: + rotary_pos_emb = rotary_pos_emb[position_ids] + else: + rotary_pos_emb = rotary_pos_emb[None, :seq_length] + rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous() + + # Run encoder. + attn_mask = full_attention_mask + if is_flash_attn_2_available: + attn_mask = attention_mask + hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder( + inputs_embeds, attn_mask, rotary_pos_emb=rotary_pos_emb, + kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states + ) + + if not return_dict: + return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=presents, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + def quantize(self, weight_bit_width: int): + from .quantization import quantize + quantize(self.encoder, weight_bit_width) + return self + + +class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel): + def __init__(self, config: ChatGLMConfig, empty_init=True, device=None): + super().__init__(config) + + self.max_sequence_length = config.max_length + self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device) + self.config = config + self.quantized = False + + if self.config.quantization_bit: + self.quantize(self.config.quantization_bit, empty_init=True) + + def _update_model_kwargs_for_generation( + self, + outputs: ModelOutput, + model_kwargs: Dict[str, Any], + is_encoder_decoder: bool = False, + standardize_cache_format: bool = False, + ) -> Dict[str, Any]: + # update past_key_values + model_kwargs["past_key_values"] = self._extract_past_from_model_output( + outputs, standardize_cache_format=standardize_cache_format + ) + + # update attention mask + if "attention_mask" in model_kwargs: + attention_mask = model_kwargs["attention_mask"] + model_kwargs["attention_mask"] = torch.cat( + [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1 + ) + + # update position ids + if "position_ids" in model_kwargs: + position_ids = model_kwargs["position_ids"] + new_position_id = position_ids[..., -1:].clone() + new_position_id += 1 + model_kwargs["position_ids"] = torch.cat( + [position_ids, new_position_id], dim=-1 + ) + + model_kwargs["is_first_forward"] = False + return model_kwargs + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + past_key_values: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + is_first_forward: bool = True, + **kwargs + ) -> dict: + # only last token for input_ids if past is not None + if position_ids is None: + position_ids = self.get_position_ids(input_ids, device=input_ids.device) + if not is_first_forward: + if past_key_values is not None: + position_ids = position_ids[..., -1:] + input_ids = input_ids[:, -1:] + return { + "input_ids": input_ids, + "past_key_values": past_key_values, + "position_ids": position_ids, + "attention_mask": attention_mask, + "return_last_logit": True, + "use_cache": use_cache + } + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Tuple[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + return_last_logit: Optional[bool] = False, + ): + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + if return_last_logit: + hidden_states = hidden_states[-1:] + lm_logits = self.transformer.output_layer(hidden_states) + lm_logits = lm_logits.transpose(0, 1).contiguous() + + loss = None + if labels is not None: + lm_logits = lm_logits.to(torch.float32) + + # Shift so that tokens < n predict n + shift_logits = lm_logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss(ignore_index=-100) + loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + + lm_logits = lm_logits.to(hidden_states.dtype) + loss = loss.to(hidden_states.dtype) + + if not return_dict: + output = (lm_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=lm_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + @staticmethod + def _reorder_cache( + past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: + """ + This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or + [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct + beam_idx at every generation step. + + Output shares the same memory storage as `past`. + """ + return tuple( + ( + layer_past[0].index_select(1, beam_idx.to(layer_past[0].device)), + layer_past[1].index_select(1, beam_idx.to(layer_past[1].device)), + ) + for layer_past in past + ) + + def process_response(self, output, history): + content = "" + history = deepcopy(history) + for response in output.split("<|assistant|>"): + metadata, content = response.split("\n", maxsplit=1) + if not metadata.strip(): + content = content.strip() + history.append({"role": "assistant", "metadata": metadata, "content": content}) + content = content.replace("[[训练时间]]", "2023年") + else: + history.append({"role": "assistant", "metadata": metadata, "content": content}) + if history[0]["role"] == "system" and "tools" in history[0]: + content = "\n".join(content.split("\n")[1:-1]) + def tool_call(**kwargs): + return kwargs + parameters = eval(content) + content = {"name": metadata.strip(), "parameters": parameters} + else: + content = {"name": metadata.strip(), "content": content} + return content, history + + @torch.inference_mode() + def chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, role: str = "user", + max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None, + **kwargs): + if history is None: + history = [] + if logits_processor is None: + logits_processor = LogitsProcessorList() + logits_processor.append(InvalidScoreLogitsProcessor()) + gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p, + "temperature": temperature, "logits_processor": logits_processor, **kwargs} + inputs = tokenizer.build_chat_input(query, history=history, role=role) + inputs = inputs.to(self.device) + eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"), + tokenizer.get_command("<|observation|>")] + outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id) + outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1] + response = tokenizer.decode(outputs) + history.append({"role": role, "content": query}) + response, history = self.process_response(response, history) + return response, history + + @torch.inference_mode() + def stream_chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, role: str = "user", + past_key_values=None,max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8, + logits_processor=None, return_past_key_values=False, **kwargs): + if history is None: + history = [] + if logits_processor is None: + logits_processor = LogitsProcessorList() + logits_processor.append(InvalidScoreLogitsProcessor()) + eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"), + tokenizer.get_command("<|observation|>")] + gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p, + "temperature": temperature, "logits_processor": logits_processor, **kwargs} + if past_key_values is None: + inputs = tokenizer.build_chat_input(query, history=history, role=role) + else: + inputs = tokenizer.build_chat_input(query, role=role) + inputs = inputs.to(self.device) + if past_key_values is not None: + past_length = past_key_values[0][0].shape[0] + if self.transformer.pre_seq_len is not None: + past_length -= self.transformer.pre_seq_len + inputs.position_ids += past_length + attention_mask = inputs.attention_mask + attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1) + inputs['attention_mask'] = attention_mask + history.append({"role": role, "content": query}) + for outputs in self.stream_generate(**inputs, past_key_values=past_key_values, + eos_token_id=eos_token_id, return_past_key_values=return_past_key_values, + **gen_kwargs): + if return_past_key_values: + outputs, past_key_values = outputs + outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1] + response = tokenizer.decode(outputs) + if response and response[-1] != "�": + response, new_history = self.process_response(response, history) + if return_past_key_values: + yield response, new_history, past_key_values + else: + yield response, new_history + + @torch.inference_mode() + def stream_generate( + self, + input_ids, + generation_config: Optional[GenerationConfig] = None, + logits_processor: Optional[LogitsProcessorList] = None, + stopping_criteria: Optional[StoppingCriteriaList] = None, + prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None, + return_past_key_values=False, + **kwargs, + ): + batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1] + + if generation_config is None: + generation_config = self.generation_config + generation_config = copy.deepcopy(generation_config) + model_kwargs = generation_config.update(**kwargs) + model_kwargs["use_cache"] = generation_config.use_cache + bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id + + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None + + has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None + if has_default_max_length and generation_config.max_new_tokens is None: + warnings.warn( + f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. " + "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we" + " recommend using `max_new_tokens` to control the maximum length of the generation.", + UserWarning, + ) + elif generation_config.max_new_tokens is not None: + generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length + if not has_default_max_length: + logger.warn( + f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(=" + f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. " + "Please refer to the documentation for more information. " + "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)", + UserWarning, + ) + + if input_ids_seq_length >= generation_config.max_length: + input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" + logger.warning( + f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to" + f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider" + " increasing `max_new_tokens`." + ) + + # 2. Set generation parameters if not already defined + logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() + stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() + + logits_processor = self._get_logits_processor( + generation_config=generation_config, + input_ids_seq_length=input_ids_seq_length, + encoder_input_ids=input_ids, + prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, + logits_processor=logits_processor, + ) + + stopping_criteria = self._get_stopping_criteria( + generation_config=generation_config, stopping_criteria=stopping_criteria + ) + logits_warper = self._get_logits_warper(generation_config) + + unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1) + scores = None + while True: + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + # forward pass to get next token + outputs = self( + **model_inputs, + return_dict=True, + output_attentions=False, + output_hidden_states=False, + ) + + next_token_logits = outputs.logits[:, -1, :] + + # pre-process distribution + next_token_scores = logits_processor(input_ids, next_token_logits) + next_token_scores = logits_warper(input_ids, next_token_scores) + + # sample + probs = nn.functional.softmax(next_token_scores, dim=-1) + if generation_config.do_sample: + next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) + else: + next_tokens = torch.argmax(probs, dim=-1) + # update generated ids, model inputs, and length for next step + input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) + model_kwargs = self._update_model_kwargs_for_generation( + outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder + ) + unfinished_sequences = unfinished_sequences.mul( + next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0) + ) + if return_past_key_values: + yield input_ids, outputs.past_key_values + else: + yield input_ids + # stop when each sentence is finished, or if we exceed the maximum length + if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores): + break + + def quantize(self, bits: int, empty_init=False, device=None, **kwargs): + if bits == 0: + return + + from .quantization import quantize + + if self.quantized: + logger.info("Already quantized.") + return self + + self.quantized = True + + self.config.quantization_bit = bits + + self.transformer.encoder = quantize(self.transformer.encoder, bits, empty_init=empty_init, device=device, + **kwargs) + return self + + +class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel): + def __init__(self, config: ChatGLMConfig, empty_init=True, device=None): + super().__init__(config) + + self.num_labels = config.num_labels + self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device) + + self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=torch.half) + if config.classifier_dropout is not None: + self.dropout = nn.Dropout(config.classifier_dropout) + else: + self.dropout = None + self.config = config + + if self.config.quantization_bit: + self.quantize(self.config.quantization_bit, empty_init=True) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + full_attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + inputs_embeds: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutputWithPast]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + full_attention_mask=full_attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + pooled_hidden_states = hidden_states[-1] + if self.dropout is not None: + pooled_hidden_states = self.dropout(pooled_hidden_states) + logits = self.classifier_head(pooled_hidden_states) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze().float(), labels.squeeze()) + else: + loss = loss_fct(logits.float(), labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels).float(), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits.float(), labels.view(-1, self.num_labels)) + + if not return_dict: + output = (logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) diff --git a/ixformer_sdk/train/speedformer/models/gpt2/__init__.py b/ixformer_sdk/train/speedformer/models/gpt2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/models/gpt2/configuration_gpt2.py b/ixformer_sdk/train/speedformer/models/gpt2/configuration_gpt2.py new file mode 100644 index 00000000..b62686fc --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/gpt2/configuration_gpt2.py @@ -0,0 +1,269 @@ +# coding=utf-8 +# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# 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. +""" OpenAI GPT-2 configuration""" +from collections import OrderedDict +from typing import Any, List, Mapping, Optional + +from transformers import PreTrainedTokenizer, TensorType, is_torch_available +from transformers.configuration_utils import PretrainedConfig +from transformers.onnx import OnnxConfigWithPast, PatchingSpec +from transformers.utils import logging + + +logger = logging.get_logger(__name__) + + +class GPT2Config(PretrainedConfig): + """ + This is the configuration class to store the configuration of a [`GPT2Model`] or a [`TFGPT2Model`]. It is used to + instantiate a GPT-2 model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the GPT-2 + [openai-community/gpt2](https://huggingface.co/openai-community/gpt2) architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 50257): + Vocabulary size of the GPT-2 model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`GPT2Model`] or [`TFGPT2Model`]. + n_positions (`int`, *optional*, defaults to 1024): + The maximum sequence length that this model might ever be used with. Typically set this to something large + just in case (e.g., 512 or 1024 or 2048). + n_embd (`int`, *optional*, defaults to 768): + Dimensionality of the embeddings and hidden states. + n_layer (`int`, *optional*, defaults to 12): + Number of hidden layers in the Transformer encoder. + n_head (`int`, *optional*, defaults to 12): + Number of attention heads for each attention layer in the Transformer encoder. + n_inner (`int`, *optional*): + Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd + activation_function (`str`, *optional*, defaults to `"gelu_new"`): + Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new"]`. + resid_pdrop (`float`, *optional*, defaults to 0.1): + The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. + embd_pdrop (`float`, *optional*, defaults to 0.1): + The dropout ratio for the embeddings. + attn_pdrop (`float`, *optional*, defaults to 0.1): + The dropout ratio for the attention. + layer_norm_epsilon (`float`, *optional*, defaults to 1e-05): + The epsilon to use in the layer normalization layers. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + summary_type (`string`, *optional*, defaults to `"cls_index"`): + Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and + [`TFGPT2DoubleHeadsModel`]. + + Has to be one of the following options: + + - `"last"`: Take the last token hidden state (like XLNet). + - `"first"`: Take the first token hidden state (like BERT). + - `"mean"`: Take the mean of all tokens hidden states. + - `"cls_index"`: Supply a Tensor of classification token position (like GPT/GPT-2). + - `"attn"`: Not implemented now, use multi-head attention. + summary_use_proj (`bool`, *optional*, defaults to `True`): + Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and + [`TFGPT2DoubleHeadsModel`]. + + Whether or not to add a projection after the vector extraction. + summary_activation (`str`, *optional*): + Argument used when doing sequence summary. Used in for the multiple choice head in + [`GPT2DoubleHeadsModel`]. + + Pass `"tanh"` for a tanh activation to the output, any other value will result in no activation. + summary_proj_to_labels (`bool`, *optional*, defaults to `True`): + Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and + [`TFGPT2DoubleHeadsModel`]. + + Whether the projection outputs should have `config.num_labels` or `config.hidden_size` classes. + summary_first_dropout (`float`, *optional*, defaults to 0.1): + Argument used when doing sequence summary, used in the models [`GPT2DoubleHeadsModel`] and + [`TFGPT2DoubleHeadsModel`]. + + The dropout ratio to be used after the projection and activation. + scale_attn_weights (`bool`, *optional*, defaults to `True`): + Scale attention weights by dividing by sqrt(hidden_size).. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). + bos_token_id (`int`, *optional*, defaults to 50256): + Id of the beginning of sentence token in the vocabulary. + eos_token_id (`int`, *optional*, defaults to 50256): + Id of the end of sentence token in the vocabulary. + scale_attn_by_inverse_layer_idx (`bool`, *optional*, defaults to `False`): + Whether to additionally scale attention weights by `1 / layer_idx + 1`. + reorder_and_upcast_attn (`bool`, *optional*, defaults to `False`): + Whether to scale keys (K) prior to computing attention (dot-product) and upcast attention + dot-product/softmax to float() when training with mixed precision. + + Example: + + ```python + >>> from transformers import GPT2Config, GPT2Model + + >>> # Initializing a GPT2 configuration + >>> configuration = GPT2Config() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = GPT2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "gpt2" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "hidden_size": "n_embd", + "max_position_embeddings": "n_positions", + "num_attention_heads": "n_head", + "num_hidden_layers": "n_layer", + } + + def __init__( + self, + vocab_size=50257, + n_positions=1024, + n_embd=768, + n_layer=12, + n_head=12, + n_inner=None, + activation_function="gelu_new", + resid_pdrop=0.1, + embd_pdrop=0.1, + attn_pdrop=0.1, + layer_norm_epsilon=1e-5, + initializer_range=0.02, + summary_type="cls_index", + summary_use_proj=True, + summary_activation=None, + summary_proj_to_labels=True, + summary_first_dropout=0.1, + scale_attn_weights=True, + use_cache=True, + bos_token_id=50256, + eos_token_id=50256, + scale_attn_by_inverse_layer_idx=False, + reorder_and_upcast_attn=False, + **kwargs, + ): + self.vocab_size = vocab_size + self.n_positions = n_positions + self.n_embd = n_embd + self.n_layer = n_layer + self.n_head = n_head + self.n_inner = n_inner + self.activation_function = activation_function + self.resid_pdrop = resid_pdrop + self.embd_pdrop = embd_pdrop + self.attn_pdrop = attn_pdrop + self.layer_norm_epsilon = layer_norm_epsilon + self.initializer_range = initializer_range + self.summary_type = summary_type + self.summary_use_proj = summary_use_proj + self.summary_activation = summary_activation + self.summary_first_dropout = summary_first_dropout + self.summary_proj_to_labels = summary_proj_to_labels + self.scale_attn_weights = scale_attn_weights + self.use_cache = use_cache + self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx + self.reorder_and_upcast_attn = reorder_and_upcast_attn + + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + + super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) + + +class GPT2OnnxConfig(OnnxConfigWithPast): + def __init__( + self, + config: PretrainedConfig, + task: str = "default", + patching_specs: List[PatchingSpec] = None, + use_past: bool = False, + ): + super().__init__(config, task=task, patching_specs=patching_specs, use_past=use_past) + if not getattr(self._config, "pad_token_id", None): + # TODO: how to do that better? + self._config.pad_token_id = 0 + + @property + def inputs(self) -> Mapping[str, Mapping[int, str]]: + common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}) + if self.use_past: + self.fill_with_past_key_values_(common_inputs, direction="inputs") + common_inputs["attention_mask"] = {0: "batch", 1: "past_sequence + sequence"} + else: + common_inputs["attention_mask"] = {0: "batch", 1: "sequence"} + + return common_inputs + + @property + def num_layers(self) -> int: + return self._config.n_layer + + @property + def num_attention_heads(self) -> int: + return self._config.n_head + + def generate_dummy_inputs( + self, + tokenizer: PreTrainedTokenizer, + batch_size: int = -1, + seq_length: int = -1, + is_pair: bool = False, + framework: Optional[TensorType] = None, + ) -> Mapping[str, Any]: + common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs( + tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework + ) + + # We need to order the input in the way they appears in the forward() + ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]}) + + # Need to add the past_keys + if self.use_past: + if not is_torch_available(): + raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.") + else: + import torch + + batch, seqlen = common_inputs["input_ids"].shape + # Not using the same length for past_key_values + past_key_values_length = seqlen + 2 + past_shape = ( + batch, + self.num_attention_heads, + past_key_values_length, + self._config.hidden_size // self.num_attention_heads, + ) + ordered_inputs["past_key_values"] = [ + (torch.zeros(past_shape), torch.zeros(past_shape)) for _ in range(self.num_layers) + ] + + ordered_inputs["attention_mask"] = common_inputs["attention_mask"] + if self.use_past: + mask_dtype = ordered_inputs["attention_mask"].dtype + ordered_inputs["attention_mask"] = torch.cat( + [ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1 + ) + + return ordered_inputs + + @property + def default_onnx_opset(self) -> int: + return 13 \ No newline at end of file diff --git a/ixformer_sdk/train/speedformer/models/gpt2/modeling_attn_mask_utils.py b/ixformer_sdk/train/speedformer/models/gpt2/modeling_attn_mask_utils.py new file mode 100644 index 00000000..67555239 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/gpt2/modeling_attn_mask_utils.py @@ -0,0 +1,500 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch + + +@dataclass +class AttentionMaskConverter: + """ + A utility attention mask class that allows one to: + - Create a causal 4d mask + - Create a causal 4d mask with slided window + - Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length, + key_value_length) that can be multiplied with attention scores + + Examples: + + ```python + >>> import torch + >>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter + + >>> converter = AttentionMaskConverter(True) + >>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32) + tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]]) + ``` + + Parameters: + is_causal (`bool`): + Whether the attention mask should be a uni-directional (causal) or bi-directional mask. + + sliding_window (`int`, *optional*): + Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer. + """ + + is_causal: bool + sliding_window: int + + def __init__(self, is_causal: bool, sliding_window: Optional[int] = None): + self.is_causal = is_causal + self.sliding_window = sliding_window + + if self.sliding_window is not None and self.sliding_window <= 0: + raise ValueError( + f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`" + ) + + def to_causal_4d( + self, + batch_size: int, + query_length: int, + key_value_length: int, + dtype: torch.dtype, + device: Union[torch.device, "str"] = "cpu", + ) -> Optional[torch.Tensor]: + """ + Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative + bias to upper right hand triangular matrix (causal mask). + """ + if not self.is_causal: + raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.") + + # If shape is not cached, create a new causal mask and cache it + input_shape = (batch_size, query_length) + past_key_values_length = key_value_length - query_length + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if input_shape[-1] > 1 or self.sliding_window is not None: + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + + return causal_4d_mask + + def to_4d( + self, + attention_mask_2d: torch.Tensor, + query_length: int, + dtype: torch.dtype, + key_value_length: Optional[int] = None, + ) -> torch.Tensor: + """ + Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length, + key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is + causal, a causal mask will be added. + """ + input_shape = (attention_mask_2d.shape[0], query_length) + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal: + if key_value_length is None: + raise ValueError( + "This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask." + ) + + past_key_values_length = key_value_length - query_length + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=attention_mask_2d.device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + elif self.sliding_window is not None: + raise NotImplementedError("Sliding window is currently only implemented for causal masking") + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to( + attention_mask_2d.device + ) + + if causal_4d_mask is not None: + expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min) + + # expanded_attn_mask + causal_4d_mask can cause some overflow + expanded_4d_mask = expanded_attn_mask + + return expanded_4d_mask + + @staticmethod + def _make_causal_mask( + input_ids_shape: torch.Size, + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, + ): + """ + Make causal mask used for bi-directional self-attention. + """ + bsz, tgt_len = input_ids_shape + mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device) + mask_cond = torch.arange(mask.size(-1), device=device) + mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0) + + mask = mask.to(dtype) + + if past_key_values_length > 0: + mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1) + + # add lower triangular sliding window mask if necessary + if sliding_window is not None: + diagonal = past_key_values_length - sliding_window + 1 + + context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal) + mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min) + + return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) + + @staticmethod + def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) + + @staticmethod + def _unmask_unattended( + expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float] + ): + # fmt: off + """ + Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when + using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + Details: https://github.com/pytorch/pytorch/issues/110213 + + `expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len]. + `attention_mask` is [bsz, src_seq_len]. + + The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias. + + For example, if `attention_mask` is + ``` + [[0, 0, 1], + [1, 1, 1], + [0, 1, 1]] + ``` + and `expanded_mask` is (e.g. here left-padding case) + ``` + [[[[0, 0, 0], + [0, 0, 0], + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[0, 0, 0], + [0, 1, 0], + [0, 1, 1]]]] + ``` + then the modified `expanded_mask` will be + ``` + [[[[1, 1, 1], <-- modified + [1, 1, 1], <-- modified + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[1, 1, 1], <-- modified + [0, 1, 0], + [0, 1, 1]]]] + ``` + """ + # fmt: on + + # Get the index of the first non-zero value for every sample in the batch. + # In the above example, indices = [[2], [0], [1]]] + tmp = torch.arange(attention_mask.shape[1], 0, -1) + indices = torch.argmax(attention_mask.cpu() * tmp, 1, keepdim=True) + + # Find the batch indexes that have unattended tokens on the leftmost side (e.g. [0, 0, 1, 1, 1]), for which the first rows of the + # expanded mask will be completely unattended. + left_masked_rows = torch.where(indices > 0)[0] + + if left_masked_rows.shape[0] == 0: + return expanded_mask + indices = indices[left_masked_rows] + + max_len = torch.max(indices) + range_tensor = torch.arange(max_len).unsqueeze(0) + range_tensor = range_tensor.repeat(indices.size(0), 1) + + # Avoid unmasking tokens at relevant target positions (on the row axis), by rather unmasking possibly several times the first row that should always be unmasked as we filtered out the batch above. + range_tensor[range_tensor >= indices] = 0 + + # TODO: we may drop support for 3D attention mask as the refactor from Patrick maybe dropped this case + if expanded_mask.dim() == 4: + num_masks = expanded_mask.shape[1] + if num_masks == 1: + # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len] + mask_slice = (left_masked_rows[:, None], 0, range_tensor) + else: + # Broadcast [left_masked_rows, 1, 1], [1, num_masks, 1], [left_masked_rows, 1, max_len] + mask_slice = ( + left_masked_rows[:, None, None], + torch.arange(num_masks)[None, :, None], + range_tensor[:, None, :], + ) + else: + # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len] + mask_slice = (left_masked_rows[:, None], range_tensor) + + expanded_mask[mask_slice] = unmasked_value + + return expanded_mask + + +def _prepare_4d_causal_attention_mask( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + attention_mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + inputs_embeds (`torch.Tensor`): + The embedded inputs as a torch Tensor. + past_key_values_length (`int`): + The length of the key value cache. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + + # 4d mask is passed through the layers + if attention_mask is not None and len(attention_mask.shape) == 2: + attention_mask = attn_mask_converter.to_4d( + attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype + ) + elif attention_mask is not None and len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + else: + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + + return attention_mask + + +# Adapted from _prepare_4d_causal_attention_mask +def _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """ + Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`. + + In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and + `key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks, + allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed). + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + batch_size, query_length = input_shape + + # torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: Fix this as well when using torchdynamo with fullgraph=True. + is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy) + + if attention_mask is not None: + # 4d mask is passed through + if len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype) + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + return attention_mask + + elif not is_tracing and torch.all(attention_mask == 1): + if query_length == 1: + # For query_length == 1, causal attention and bi-directional attention are the same. + attention_mask = None + elif key_value_length == query_length: + attention_mask = None + else: + # Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation + # may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + pass + elif query_length > 1 and key_value_length != query_length: + # See the comment above (https://github.com/pytorch/pytorch/issues/108108). + # Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`. + attention_mask = True + elif is_tracing: + raise ValueError( + 'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.' + ) + + if attention_mask is None: + expanded_4d_mask = None + elif attention_mask is True: + expanded_4d_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + else: + expanded_4d_mask = attn_mask_converter.to_4d( + attention_mask, + input_shape[-1], + dtype=inputs_embeds.dtype, + key_value_length=key_value_length, + ) + + # From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend + # produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213 + # + # This fix is not applied in case we are tracing with torch.jit.trace or symbolic_trace, as _unmask_unattended has a data-dependent + # controlflow that can not be captured properly. + # TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case. + if query_length > 1 and not is_tracing: + expanded_4d_mask = AttentionMaskConverter._unmask_unattended( + expanded_4d_mask, attention_mask, unmasked_value=0.0 + ) + + return expanded_4d_mask + + +def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + batch_size, key_value_length = mask.shape + tgt_len = tgt_len if tgt_len is not None else key_value_length + + # torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: Fix this as well when using torchdynamo with fullgraph=True. + is_tracing = torch.jit.is_tracing() + + if torch.all(mask == 1): + if is_tracing: + pass + elif tgt_len == 1: + # For query_length == 1, causal attention and bi-directional attention are the same. + return None + elif key_value_length == tgt_len: + return None + else: + # Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation + # may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + else: + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _create_4d_causal_attention_mask( + input_shape: Union[torch.Size, Tuple, List], + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, +) -> Optional[torch.Tensor]: + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` + + Args: + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + device (`int`): + The torch device the created mask shall have. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = past_key_values_length + input_shape[-1] + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device + ) + + return attention_mask diff --git a/ixformer_sdk/train/speedformer/models/gpt2/modeling_gpt2.py b/ixformer_sdk/train/speedformer/models/gpt2/modeling_gpt2.py new file mode 100644 index 00000000..c7aaf752 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/gpt2/modeling_gpt2.py @@ -0,0 +1,1949 @@ +# coding=utf-8 +# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# 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. +"""PyTorch OpenAI GPT-2 model.""" + +import math +import os +import warnings +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.cuda.amp import autocast +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.modeling_outputs import ( + BaseModelOutputWithPastAndCrossAttentions, + CausalLMOutputWithCrossAttentions, + QuestionAnsweringModelOutput, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from transformers.modeling_utils import PreTrainedModel, SequenceSummary +from transformers.pytorch_utils import Conv1D, find_pruneable_heads_and_indices, prune_conv1d_layer +from transformers.utils import ( + ModelOutput, + add_code_sample_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from transformers.utils.model_parallel_utils import assert_device_map, get_device_map +from .configuration_gpt2 import GPT2Config + + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input + + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "openai-community/gpt2" +_CONFIG_FOR_DOC = "GPT2Config" + + +# Copied from transformers.models.llama.modeling_llama._get_unpad_data +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path): + """Load tf checkpoints in a pytorch model""" + try: + import re + + import tensorflow as tf + except ImportError: + logger.error( + "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " + "https://www.tensorflow.org/install/ for installation instructions." + ) + raise + tf_path = os.path.abspath(gpt2_checkpoint_path) + logger.info(f"Converting TensorFlow checkpoint from {tf_path}") + # Load weights from TF model + init_vars = tf.train.list_variables(tf_path) + names = [] + arrays = [] + for name, shape in init_vars: + logger.info(f"Loading TF weight {name} with shape {shape}") + array = tf.train.load_variable(tf_path, name) + names.append(name) + arrays.append(array.squeeze()) + + for name, array in zip(names, arrays): + name = name[6:] # skip "model/" + name = name.split("/") + pointer = model + for m_name in name: + if re.fullmatch(r"[A-Za-z]+\d+", m_name): + scope_names = re.split(r"(\d+)", m_name) + else: + scope_names = [m_name] + if scope_names[0] == "w" or scope_names[0] == "g": + pointer = getattr(pointer, "weight") + elif scope_names[0] == "b": + pointer = getattr(pointer, "bias") + elif scope_names[0] == "wpe" or scope_names[0] == "wte": + pointer = getattr(pointer, scope_names[0]) + pointer = getattr(pointer, "weight") + else: + pointer = getattr(pointer, scope_names[0]) + if len(scope_names) >= 2: + num = int(scope_names[1]) + pointer = pointer[num] + try: + if pointer.shape != array.shape: + raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched") + except ValueError as e: + e.args += (pointer.shape, array.shape) + raise + logger.info(f"Initialize PyTorch weight {name}") + pointer.data = torch.from_numpy(array) + return model + + +class GPT2Attention(nn.Module): + def __init__(self, config, is_cross_attention=False, layer_idx=None): + super().__init__() + self.config = config + max_positions = config.max_position_embeddings + self.register_buffer( + "bias", + torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view( + 1, 1, max_positions, max_positions + ), + persistent=False, + ) + self.register_buffer("masked_bias", torch.tensor(-1e4), persistent=False) + + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + self.split_size = self.embed_dim + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" + f" {self.num_heads})." + ) + + self.scale_attn_weights = config.scale_attn_weights + self.is_cross_attention = is_cross_attention + + # Layer-wise attention scaling, reordering, and upcasting + self.scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx + self.layer_idx = layer_idx + self.reorder_and_upcast_attn = config.reorder_and_upcast_attn + + if self.is_cross_attention: + self.c_attn = Conv1D(2 * self.embed_dim, self.embed_dim) + self.q_attn = Conv1D(self.embed_dim, self.embed_dim) + else: + self.c_attn = Conv1D(3 * self.embed_dim, self.embed_dim) + self.c_proj = Conv1D(self.embed_dim, self.embed_dim) + + self.attn_dropout = nn.Dropout(config.attn_pdrop) + self.resid_dropout = nn.Dropout(config.resid_pdrop) + self.is_causal = True + + self.pruned_heads = set() + + def prune_heads(self, heads): + if len(heads) == 0: + return + heads, index = find_pruneable_heads_and_indices(heads, self.num_heads, self.head_dim, self.pruned_heads) + index_attn = torch.cat([index, index + self.split_size, index + (2 * self.split_size)]) + + # Prune conv1d layers + self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1) + self.c_proj = prune_conv1d_layer(self.c_proj, index, dim=0) + + # Update hyper params + self.split_size = (self.split_size // self.num_heads) * (self.num_heads - len(heads)) + self.num_heads = self.num_heads - len(heads) + self.pruned_heads = self.pruned_heads.union(heads) + + def _attn(self, query, key, value, attention_mask=None, head_mask=None): + attn_weights = torch.matmul(query, key.transpose(-1, -2)) + + if self.scale_attn_weights: + attn_weights = attn_weights / torch.full( + [], value.size(-1) ** 0.5, dtype=attn_weights.dtype, device=attn_weights.device + ) + + # Layer-wise attention scaling + if self.scale_attn_by_inverse_layer_idx: + attn_weights = attn_weights / float(self.layer_idx + 1) + + if not self.is_cross_attention: + # if only "normal" attention layer implements causal mask + query_length, key_length = query.size(-2), key.size(-2) + causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length] + mask_value = torch.finfo(attn_weights.dtype).min + # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`. + # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device` + mask_value = torch.full([], mask_value, dtype=attn_weights.dtype, device=attn_weights.device) + attn_weights = torch.where(causal_mask, attn_weights.to(attn_weights.dtype), mask_value) + + if attention_mask is not None: + # Apply the attention mask + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op otherwise + attn_weights = attn_weights.type(value.dtype) + attn_weights = self.attn_dropout(attn_weights) + + # Mask heads if we want to + if head_mask is not None: + attn_weights = attn_weights * head_mask + + attn_output = torch.matmul(attn_weights, value) + + return attn_output, attn_weights + + def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, head_mask=None): + # Use `torch.baddbmm` (a bit more efficient w/ alpha param for scaling -- from Megatron-LM) + bsz, num_heads, q_seq_len, dk = query.size() + _, _, k_seq_len, _ = key.size() + + # Preallocate attn_weights for `baddbmm` + attn_weights = torch.empty(bsz * num_heads, q_seq_len, k_seq_len, dtype=torch.float32, device=query.device) + + # Compute Scale Factor + scale_factor = 1.0 + if self.scale_attn_weights: + scale_factor /= float(value.size(-1)) ** 0.5 + + if self.scale_attn_by_inverse_layer_idx: + scale_factor /= float(self.layer_idx + 1) + + # Upcast (turn off autocast) and reorder (Scale K by 1 / root(dk)) + with autocast(enabled=False): + q, k = query.reshape(-1, q_seq_len, dk), key.transpose(-1, -2).reshape(-1, dk, k_seq_len) + attn_weights = torch.baddbmm(attn_weights, q.float(), k.float(), beta=0, alpha=scale_factor) + attn_weights = attn_weights.reshape(bsz, num_heads, q_seq_len, k_seq_len) + + if not self.is_cross_attention: + # if only "normal" attention layer implements causal mask + query_length, key_length = query.size(-2), key.size(-2) + causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length] + mask_value = torch.finfo(attn_weights.dtype).min + # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`. + # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device` + mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device) + attn_weights = torch.where(causal_mask, attn_weights, mask_value) + + if attention_mask is not None: + # Apply the attention mask + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op if otherwise + if attn_weights.dtype != torch.float32: + raise RuntimeError("Error with upcasting, attn_weights does not have dtype torch.float32") + attn_weights = attn_weights.type(value.dtype) + attn_weights = self.attn_dropout(attn_weights) + + # Mask heads if we want to + if head_mask is not None: + attn_weights = attn_weights * head_mask + + attn_output = torch.matmul(attn_weights, value) + + return attn_output, attn_weights + + def _split_heads(self, tensor, num_heads, attn_head_size): + """ + Splits hidden_size dim into attn_head_size and num_heads + """ + new_shape = tensor.size()[:-1] + (num_heads, attn_head_size) + tensor = tensor.view(new_shape) + return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features) + + def _merge_heads(self, tensor, num_heads, attn_head_size): + """ + Merges attn_head_size dim and num_attn_heads dim into hidden_size + """ + tensor = tensor.permute(0, 2, 1, 3).contiguous() + new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,) + return tensor.view(new_shape) + + def forward( + self, + hidden_states: Optional[Tuple[torch.FloatTensor]], + layer_past: Optional[Tuple[torch.Tensor]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]: + if encoder_hidden_states is not None: + if not hasattr(self, "q_attn"): + raise ValueError( + "If class is used as cross attention, the weights `q_attn` have to be defined. " + "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`." + ) + + query = self.q_attn(hidden_states) + key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2) + attention_mask = encoder_attention_mask + else: + query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2) + + query = self._split_heads(query, self.num_heads, self.head_dim) + key = self._split_heads(key, self.num_heads, self.head_dim) + value = self._split_heads(value, self.num_heads, self.head_dim) + + if layer_past is not None: + past_key, past_value = layer_past + key = torch.cat((past_key, key), dim=-2) + value = torch.cat((past_value, value), dim=-2) + + if use_cache is True: + present = (key, value) + else: + present = None + + if self.reorder_and_upcast_attn: + attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask) + else: + attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask) + + attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim) + attn_output = self.c_proj(attn_output) + attn_output = self.resid_dropout(attn_output) + + outputs = (attn_output, present) + if output_attentions: + outputs += (attn_weights,) + + return outputs # a, present, (attentions) + + +class GPT2FlashAttention2(GPT2Attention): + """ + GPT2 flash attention module. This module inherits from `GPT2Attention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: Optional[Tuple[torch.FloatTensor]], + layer_past: Optional[Tuple[torch.Tensor]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]: + bsz, _, _ = hidden_states.size() + if encoder_hidden_states is not None: + if not hasattr(self, "q_attn"): + raise ValueError( + "If class is used as cross attention, the weights `q_attn` have to be defined. " + "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`." + ) + + query = self.q_attn(hidden_states) + key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2) + attention_mask = encoder_attention_mask + else: + query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2) + + query = self._split_heads(query, self.num_heads, self.head_dim) + key = self._split_heads(key, self.num_heads, self.head_dim) + value = self._split_heads(value, self.num_heads, self.head_dim) + + if layer_past is not None: + past_key = layer_past[0] + past_value = layer_past[1] + key = torch.cat((past_key, key), dim=-2) + value = torch.cat((past_value, value), dim=-2) + + present = None + if use_cache is True: + present = (key, value) + + query_length = query.shape[2] + tgt_len = key.shape[2] + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + query = query.transpose(1, 2).view(bsz, query_length, self.num_heads, self.head_dim) + key = key.transpose(1, 2).view(bsz, tgt_len, self.num_heads, self.head_dim) + value = value.transpose(1, 2).view(bsz, tgt_len, self.num_heads, self.head_dim) + + attn_dropout = self.attn_dropout.p if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (LlamaRMSNorm handles it correctly) + + if query.dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.c_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query = query.to(target_dtype) + key = key.to(target_dtype) + value = value.to(target_dtype) + + attn_output = self._flash_attention_forward( + query, key, value, attention_mask, query_length, dropout=attn_dropout + ) + + attn_weights_reshaped = attn_output.reshape(bsz, query_length, self.num_heads * self.head_dim) + attn_output = self.c_proj(attn_weights_reshaped) + attn_output = self.resid_dropout(attn_output) + + outputs = (attn_output, present) + if output_attentions: + outputs += (attn_weights_reshaped,) + + return outputs + + # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward + def _flash_attention_forward( + self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None + ): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`float`): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + """ + if not self._flash_attn_uses_top_left_mask: + causal = self.is_causal + else: + # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__. + causal = self.is_causal and query_length != 1 + + # Contains at least one padding token in the sequence + if attention_mask is not None: + batch_size = query_states.shape[0] + query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( + query_states, key_states, value_states, attention_mask, query_length + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) + else: + attn_output = flash_attn_func( + query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal + ) + + return attn_output + + # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input + def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = index_first_axis( + key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +class GPT2MLP(nn.Module): + def __init__(self, intermediate_size, config): + super().__init__() + embed_dim = config.hidden_size + self.c_fc = Conv1D(intermediate_size, embed_dim) + self.c_proj = Conv1D(embed_dim, intermediate_size) + self.act = ACT2FN[config.activation_function] + self.dropout = nn.Dropout(config.resid_pdrop) + + def forward(self, hidden_states: Optional[Tuple[torch.FloatTensor]]) -> torch.FloatTensor: + hidden_states = self.c_fc(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.c_proj(hidden_states) + hidden_states = self.dropout(hidden_states) + return hidden_states + + +GPT2_ATTENTION_CLASSES = { + "eager": GPT2Attention, + "flash_attention_2": GPT2FlashAttention2, +} + + +class GPT2Block(nn.Module): + def __init__(self, config, layer_idx=None): + super().__init__() + hidden_size = config.hidden_size + inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size + attention_class = GPT2_ATTENTION_CLASSES[config._attn_implementation] + + self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + self.attn = attention_class(config=config, layer_idx=layer_idx) + self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + + if config.add_cross_attention: + self.crossattention = attention_class(config=config, is_cross_attention=True, layer_idx=layer_idx) + self.ln_cross_attn = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + + self.mlp = GPT2MLP(inner_dim, config) + + def forward( + self, + hidden_states: Optional[Tuple[torch.FloatTensor]], + layer_past: Optional[Tuple[torch.Tensor]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]: + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attn_outputs = self.attn( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + head_mask=head_mask, + use_cache=use_cache, + output_attentions=output_attentions, + ) + attn_output = attn_outputs[0] # output_attn: a, present, (attentions) + outputs = attn_outputs[1:] + # residual connection + hidden_states = attn_output + residual + + if encoder_hidden_states is not None: + # add one self-attention block for cross-attention + if not hasattr(self, "crossattention"): + raise ValueError( + f"If `encoder_hidden_states` are passed, {self} has to be instantiated with " + "cross-attention layers by setting `config.add_cross_attention=True`" + ) + residual = hidden_states + hidden_states = self.ln_cross_attn(hidden_states) + cross_attn_outputs = self.crossattention( + hidden_states, + attention_mask=attention_mask, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + ) + attn_output = cross_attn_outputs[0] + # residual connection + hidden_states = residual + attn_output + outputs = outputs + cross_attn_outputs[2:] # add cross attentions if we output attention weights + + residual = hidden_states + hidden_states = self.ln_2(hidden_states) + feed_forward_hidden_states = self.mlp(hidden_states) + # residual connection + hidden_states = residual + feed_forward_hidden_states + + if use_cache: + outputs = (hidden_states,) + outputs + else: + outputs = (hidden_states,) + outputs[1:] + + return outputs # hidden_states, present, (attentions, cross_attentions) + + +class GPT2PreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = GPT2Config + load_tf_weights = load_tf_weights_in_gpt2 + base_model_prefix = "transformer" + is_parallelizable = True + supports_gradient_checkpointing = True + _no_split_modules = ["GPT2Block"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights(self, module): + """Initialize the weights.""" + if isinstance(module, (nn.Linear, Conv1D)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + for name, p in module.named_parameters(): + if name == "c_proj.weight": + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + p.data.normal_(mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.n_layer))) + + +@dataclass +class GPT2DoubleHeadsModelOutput(ModelOutput): + """ + Base class for outputs of models predicting if two sentences are consecutive or not. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss. + mc_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mc_labels` is provided): + Multiple choice classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, num_choices, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + mc_logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`): + Prediction scores of the multiple choice classification head (scores for each choice before SoftMax). + past_key_values (`Tuple[Tuple[torch.Tensor]]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of length `config.n_layers`, containing tuples of tensors of shape `(batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of + shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + GPT2Attentions weights after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + loss: Optional[torch.FloatTensor] = None + mc_loss: Optional[torch.FloatTensor] = None + logits: torch.FloatTensor = None + mc_logits: torch.FloatTensor = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + + +GPT2_START_DOCSTRING = r""" + + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`GPT2Config`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +GPT2_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): + `input_ids_length` = `sequence_length` if `past_key_values` is `None` else + `past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input + sequence tokens in the vocabulary. + + If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as + `input_ids`. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.n_layers`): + Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see + `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have + their past given to this model should not be passed as `input_ids` as they have already been computed. + attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + If `past_key_values` is used, `attention_mask` needs to contain the masking strategy that was used for + `past_key_values`. In other words, the `attention_mask` always has to have the length: + `len(past_key_values) + len(input_ids)` + + [What are attention masks?](../glossary#attention-mask) + token_type_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*): + Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, + 1]`: + + - 0 corresponds to a *sentence A* token, + - 1 corresponds to a *sentence B* token. + + [What are token type IDs?](../glossary#token-type-ids) + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.max_position_embeddings - 1]`. + + [What are position IDs?](../glossary#position-ids) + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + + If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see + `past_key_values`). + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" +PARALLELIZE_DOCSTRING = r""" + This is an experimental feature and is a subject to change at a moment's notice. + + Uses a device map to distribute attention modules of the model across several devices. If no device map is given, + it will evenly distribute blocks across all devices. + + Args: + device_map (`Dict[int, list]`, optional, defaults to None): + A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always + automatically mapped to the first device (for esoteric reasons). That means that the first device should + have fewer attention modules mapped to it than other devices. For reference, the gpt2 models have the + following number of attention modules: + + - openai-community/gpt2: 12 + - openai-community/gpt2-medium: 24 + - openai-community/gpt2-large: 36 + - openai-community/gpt2-xl: 48 + + Example: + + ```python + # Here is an example of a device map on a machine with 4 GPUs using gpt2-xl, which has a total of 48 attention modules: + model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2-xl") + device_map = { + 0: [0, 1, 2, 3, 4, 5, 6, 7, 8], + 1: [9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21], + 2: [22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], + 3: [35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47], + } + model.parallelize(device_map) + ``` +""" +DEPARALLELIZE_DOCSTRING = r""" + Moves the model to cpu from a model parallel state. + + Example: + + ```python + # On a 4 GPU machine with openai-community/gpt2-large: + model = GPT2LMHeadModel.from_pretrained("openai-community/gpt2-large") + device_map = { + 0: [0, 1, 2, 3, 4, 5, 6, 7], + 1: [8, 9, 10, 11, 12, 13, 14, 15], + 2: [16, 17, 18, 19, 20, 21, 22, 23], + 3: [24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35], + } + model.parallelize(device_map) # Splits the model across several devices + model.deparallelize() # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache() + ``` +""" + + +@add_start_docstrings( + "The bare GPT2 Model transformer outputting raw hidden-states without any specific head on top.", + GPT2_START_DOCSTRING, +) +class GPT2Model(GPT2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embed_dim = config.hidden_size + + self.wte = nn.Embedding(config.vocab_size, self.embed_dim) + self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim) + + self.drop = nn.Dropout(config.embd_pdrop) + self.h = nn.ModuleList([GPT2Block(config, layer_idx=i) for i in range(config.num_hidden_layers)]) + self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) + + # Model parallel + self.model_parallel = False + self.device_map = None + self.gradient_checkpointing = False + self._attn_implementation = config._attn_implementation + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings(PARALLELIZE_DOCSTRING) + def parallelize(self, device_map=None): + # Check validity of device_map + warnings.warn( + "`GPT2Model.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your" + " model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" + " `device_map` but it needs to be a dictionary module_name to device, so for instance {'h.0': 0, 'h.1': 1," + " ...}", + FutureWarning, + ) + self.device_map = ( + get_device_map(len(self.h), range(torch.cuda.device_count())) if device_map is None else device_map + ) + assert_device_map(self.device_map, len(self.h)) + self.model_parallel = True + self.first_device = "cpu" if "cpu" in self.device_map.keys() else "cuda:" + str(min(self.device_map.keys())) + self.last_device = "cuda:" + str(max(self.device_map.keys())) + self.wte = self.wte.to(self.first_device) + self.wpe = self.wpe.to(self.first_device) + # Load onto devices + for k, v in self.device_map.items(): + for block in v: + cuda_device = "cuda:" + str(k) + self.h[block] = self.h[block].to(cuda_device) + # ln_f to last + self.ln_f = self.ln_f.to(self.last_device) + + @add_start_docstrings(DEPARALLELIZE_DOCSTRING) + def deparallelize(self): + warnings.warn( + "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", + FutureWarning, + ) + self.model_parallel = False + self.device_map = None + self.first_device = "cpu" + self.last_device = "cpu" + self.wte = self.wte.to("cpu") + self.wpe = self.wpe.to("cpu") + for index in range(len(self.h)): + self.h[index] = self.h[index].to("cpu") + self.ln_f = self.ln_f.to("cpu") + torch.cuda.empty_cache() + + def get_input_embeddings(self): + return self.wte + + def set_input_embeddings(self, new_embeddings): + self.wte = new_embeddings + + def _prune_heads(self, heads_to_prune): + """ + Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} + """ + for layer, heads in heads_to_prune.items(): + self.h[layer].attn.prune_heads(heads) + + @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=BaseModelOutputWithPastAndCrossAttentions, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + token_type_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + batch_size = input_ids.shape[0] + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + batch_size = inputs_embeds.shape[0] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if token_type_ids is not None: + token_type_ids = token_type_ids.view(-1, input_shape[-1]) + + if past_key_values is None: + past_length = 0 + past_key_values = tuple([None] * len(self.h)) + else: + past_length = past_key_values[0][0].size(-2) + if position_ids is None: + position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device) + position_ids = position_ids.unsqueeze(0) + + # Attention mask. + if attention_mask is not None: + attention_mask = attention_mask.view(batch_size, -1) + if self._attn_implementation == "flash_attention_2": + attention_mask = attention_mask if 0 in attention_mask else None + else: + # We create a 3D attention mask from a 2D tensor mask. + # Sizes are [batch_size, 1, 1, to_seq_length] + # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length] + # this attention mask is more simple than the triangular masking of causal attention + # used in OpenAI GPT, we just need to prepare the broadcast dimension here. + attention_mask = attention_mask[:, None, None, :] + + # Since attention_mask is 1.0 for positions we want to attend and 0.0 for + # masked positions, this operation will create a tensor which is 0.0 for + # positions we want to attend and the dtype's smallest value for masked positions. + # Since we are adding it to the raw scores before the softmax, this is + # effectively the same as removing these entirely. + attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility + attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min + + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.add_cross_attention and encoder_hidden_states is not None: + encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) + if self._attn_implementation != "flash_attention_2": + encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask) + else: + encoder_attention_mask = None + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # head_mask has shape n_layer x batch x n_heads x N x N + head_mask = self.get_head_mask(head_mask, self.config.n_layer) + + if inputs_embeds is None: + inputs_embeds = self.wte(input_ids) + position_embeds = self.wpe(position_ids) + hidden_states = inputs_embeds + position_embeds + + if token_type_ids is not None: + token_type_embeds = self.wte(token_type_ids) + hidden_states = hidden_states + token_type_embeds + + hidden_states = self.drop(hidden_states) + + output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + presents = () if use_cache else None + all_self_attentions = () if output_attentions else None + all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None + all_hidden_states = () if output_hidden_states else None + for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)): + # Model parallel + if self.model_parallel: + torch.cuda.set_device(hidden_states.device) + # Ensure layer_past is on same device as hidden_states (might not be correct) + if layer_past is not None: + layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past) + # Ensure that attention_mask is always on the same device as hidden_states + if attention_mask is not None: + attention_mask = attention_mask.to(hidden_states.device) + if isinstance(head_mask, torch.Tensor): + head_mask = head_mask.to(hidden_states.device) + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if self.gradient_checkpointing and self.training: + outputs = self._gradient_checkpointing_func( + block.__call__, + hidden_states, + None, + attention_mask, + head_mask[i], + encoder_hidden_states, + encoder_attention_mask, + use_cache, + output_attentions, + ) + else: + outputs = block( + hidden_states, + layer_past=layer_past, + attention_mask=attention_mask, + head_mask=head_mask[i], + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + ) + + hidden_states = outputs[0] + if use_cache is True: + presents = presents + (outputs[1],) + + if output_attentions: + all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],) + if self.config.add_cross_attention: + all_cross_attentions = all_cross_attentions + (outputs[3 if use_cache else 2],) + + # Model Parallel: If it's the last layer for that device, put things on the next device + if self.model_parallel: + for k, v in self.device_map.items(): + if i == v[-1] and "cuda:" + str(k) != self.last_device: + hidden_states = hidden_states.to("cuda:" + str(k + 1)) + + hidden_states = self.ln_f(hidden_states) + + hidden_states = hidden_states.view(output_shape) + # Add last hidden state + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, presents, all_hidden_states, all_self_attentions, all_cross_attentions] + if v is not None + ) + + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=presents, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + cross_attentions=all_cross_attentions, + ) + + +@add_start_docstrings( + """ + The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input + embeddings). + """, + GPT2_START_DOCSTRING, +) +class GPT2LMHeadModel(GPT2PreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.transformer = GPT2Model(config) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + + # Model parallel + self.model_parallel = False + self.device_map = None + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings(PARALLELIZE_DOCSTRING) + def parallelize(self, device_map=None): + warnings.warn( + "`GPT2LMHeadModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should load" + " your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" + " `device_map` but it needs to be a dictionary module_name to device, so for instance {'transformer.h.0':" + " 0, 'transformer.h.1': 1, ...}", + FutureWarning, + ) + self.device_map = ( + get_device_map(len(self.transformer.h), range(torch.cuda.device_count())) + if device_map is None + else device_map + ) + assert_device_map(self.device_map, len(self.transformer.h)) + self.transformer.parallelize(self.device_map) + self.lm_head = self.lm_head.to(self.transformer.first_device) + self.model_parallel = True + + @add_start_docstrings(DEPARALLELIZE_DOCSTRING) + def deparallelize(self): + warnings.warn( + "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", + FutureWarning, + ) + self.transformer.deparallelize() + self.transformer = self.transformer.to("cpu") + self.lm_head = self.lm_head.to("cpu") + self.model_parallel = False + torch.cuda.empty_cache() + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs): + token_type_ids = kwargs.get("token_type_ids", None) + # Omit tokens covered by past_key_values + if past_key_values: + past_length = past_key_values[0][0].shape[2] + + # Some generation methods already pass only the last input ID + if input_ids.shape[1] > past_length: + remove_prefix_length = past_length + else: + # Default to old behavior: keep only final ID + remove_prefix_length = input_ids.shape[1] - 1 + + input_ids = input_ids[:, remove_prefix_length:] + if token_type_ids is not None: + token_type_ids = token_type_ids[:, -input_ids.shape[1] :] + + attention_mask = kwargs.get("attention_mask", None) + position_ids = kwargs.get("position_ids", None) + + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + else: + position_ids = None + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "position_ids": position_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + } + ) + + return model_inputs + + @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=CausalLMOutputWithCrossAttentions, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + token_type_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set + `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` + are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + + # Set device for model parallelism + if self.model_parallel: + torch.cuda.set_device(self.transformer.first_device) + hidden_states = hidden_states.to(self.lm_head.weight.device) + + lm_logits = self.lm_head(hidden_states) + + loss = None + if labels is not None: + # move labels to correct device to enable model parallelism + labels = labels.to(lm_logits.device) + # Shift so that tokens < n predict n + shift_logits = lm_logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + + if not return_dict: + output = (lm_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutputWithCrossAttentions( + loss=loss, + logits=lm_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + cross_attentions=transformer_outputs.cross_attentions, + ) + + @staticmethod + def _reorder_cache( + past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor + ) -> Tuple[Tuple[torch.Tensor]]: + """ + This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or + [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct + beam_idx at every generation step. + """ + return tuple( + tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past) + for layer_past in past_key_values + ) + + +@add_start_docstrings( + """ +The GPT2 Model transformer with a language modeling and a multiple-choice classification head on top e.g. for +RocStories/SWAG tasks. The two heads are two linear layers. The language modeling head has its weights tied to the +input embeddings, the classification head takes as input the input of a specified classification token index in the +input sequence). +""", + GPT2_START_DOCSTRING, +) +class GPT2DoubleHeadsModel(GPT2PreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + config.num_labels = 1 + self.transformer = GPT2Model(config) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + self.multiple_choice_head = SequenceSummary(config) + + # Model parallel + self.model_parallel = False + self.device_map = None + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings(PARALLELIZE_DOCSTRING) + def parallelize(self, device_map=None): + warnings.warn( + "`GPT2DoubleHeadsModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should" + " load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your" + " own `device_map` but it needs to be a dictionary module_name to device, so for instance" + " {'transformer.h.0': 0, 'transformer.h.1': 1, ...}", + FutureWarning, + ) + self.device_map = ( + get_device_map(len(self.transformer.h), range(torch.cuda.device_count())) + if device_map is None + else device_map + ) + assert_device_map(self.device_map, len(self.transformer.h)) + self.transformer.parallelize(self.device_map) + self.lm_head = self.lm_head.to(self.transformer.first_device) + self.multiple_choice_head = self.multiple_choice_head.to(self.transformer.first_device) + self.model_parallel = True + + @add_start_docstrings(DEPARALLELIZE_DOCSTRING) + def deparallelize(self): + warnings.warn( + "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", + FutureWarning, + ) + self.transformer.deparallelize() + self.transformer = self.transformer.to("cpu") + self.lm_head = self.lm_head.to("cpu") + self.multiple_choice_head = self.multiple_choice_head.to("cpu") + self.model_parallel = False + torch.cuda.empty_cache() + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def prepare_inputs_for_generation(self, input_ids, inputs_embeds=None, past_key_values=None, **kwargs): + token_type_ids = kwargs.get("token_type_ids", None) + # Omit tokens covered by past_key_values + if past_key_values: + past_length = past_key_values[0][0].shape[2] + + # Some generation methods already pass only the last input ID + if input_ids.shape[1] > past_length: + remove_prefix_length = past_length + else: + # Default to old behavior: keep only final ID + remove_prefix_length = input_ids.shape[1] - 1 + + input_ids = input_ids[:, remove_prefix_length:] + if token_type_ids is not None: + token_type_ids = token_type_ids[:, -input_ids.shape[1] :] + + attention_mask = kwargs.get("attention_mask", None) + position_ids = kwargs.get("position_ids", None) + + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + else: + position_ids = None + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids.contiguous()} + + model_inputs.update( + { + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "position_ids": position_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + } + ) + return model_inputs + + @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=GPT2DoubleHeadsModelOutput, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + token_type_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + mc_token_ids: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + mc_labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs, + ) -> Union[Tuple, GPT2DoubleHeadsModelOutput]: + r""" + mc_token_ids (`torch.LongTensor` of shape `(batch_size, num_choices)`, *optional*, default to index of the last token of the input): + Index of the classification token in each input sequence. Selected in the range `[0, input_ids.size(-1) - + 1]`. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set + `labels = input_ids`. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`. All labels set to + `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size - 1]` + mc_labels (`torch.LongTensor` of shape `(batch_size)`, *optional*): + Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]` + where *num_choices* is the size of the second dimension of the input tensors. (see *input_ids* above) + + Return: + + Example: + + ```python + >>> import torch + >>> from transformers import AutoTokenizer, GPT2DoubleHeadsModel + + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> model = GPT2DoubleHeadsModel.from_pretrained("openai-community/gpt2") + + >>> # Add a [CLS] to the vocabulary (we should train it also!) + >>> num_added_tokens = tokenizer.add_special_tokens({"cls_token": "[CLS]"}) + >>> # Update the model embeddings with the new vocabulary size + >>> embedding_layer = model.resize_token_embeddings(len(tokenizer)) + + >>> choices = ["Hello, my dog is cute [CLS]", "Hello, my cat is cute [CLS]"] + >>> encoded_choices = [tokenizer.encode(s) for s in choices] + >>> cls_token_location = [tokens.index(tokenizer.cls_token_id) for tokens in encoded_choices] + + >>> input_ids = torch.tensor(encoded_choices).unsqueeze(0) # Batch size: 1, number of choices: 2 + >>> mc_token_ids = torch.tensor([cls_token_location]) # Batch size: 1 + + >>> outputs = model(input_ids, mc_token_ids=mc_token_ids) + >>> lm_logits = outputs.logits + >>> mc_logits = outputs.mc_logits + ```""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + + # Set device for model parallelism + if self.model_parallel: + torch.cuda.set_device(self.transformer.first_device) + hidden_states = hidden_states.to(self.lm_head.weight.device) + + lm_logits = self.lm_head(hidden_states) + mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids).squeeze(-1) + + mc_loss = None + if mc_labels is not None: + loss_fct = CrossEntropyLoss() + mc_loss = loss_fct(mc_logits.view(-1, mc_logits.size(-1)), mc_labels.view(-1)) + lm_loss = None + if labels is not None: + labels = labels.to(lm_logits.device) + shift_logits = lm_logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + loss_fct = CrossEntropyLoss() + lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + + if not return_dict: + output = (lm_logits, mc_logits) + transformer_outputs[1:] + if mc_loss is not None: + output = (mc_loss,) + output + return ((lm_loss,) + output) if lm_loss is not None else output + + return GPT2DoubleHeadsModelOutput( + loss=lm_loss, + mc_loss=mc_loss, + logits=lm_logits, + mc_logits=mc_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + @staticmethod + def _reorder_cache( + past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor + ) -> Tuple[Tuple[torch.Tensor]]: + """ + This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or + [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct + beam_idx at every generation step. + """ + return tuple( + tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past) + for layer_past in past_key_values + ) + + +@add_start_docstrings( + """ + The GPT2 Model transformer with a sequence classification head on top (linear layer). + + [`GPT2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-1) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + GPT2_START_DOCSTRING, +) +class GPT2ForSequenceClassification(GPT2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.transformer = GPT2Model(config) + self.score = nn.Linear(config.n_embd, self.num_labels, bias=False) + + # Model parallel + self.model_parallel = False + self.device_map = None + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint="microsoft/DialogRPT-updown", + output_type=SequenceClassifierOutputWithPast, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + token_type_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size, sequence_length = input_ids.shape[:2] + else: + batch_size, sequence_length = inputs_embeds.shape[:2] + + assert ( + self.config.pad_token_id is not None or batch_size == 1 + ), "Cannot handle batch sizes > 1 if no padding token is defined." + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + logger.warning( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@add_start_docstrings( + """ + GPT2 Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for + Named-Entity-Recognition (NER) tasks. + """, + GPT2_START_DOCSTRING, +) +class GPT2ForTokenClassification(GPT2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.transformer = GPT2Model(config) + if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None: + classifier_dropout = config.classifier_dropout + elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Model parallel + self.model_parallel = False + self.device_map = None + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING) + # fmt: off + @add_code_sample_docstrings( + checkpoint="brad1141/gpt2-finetuned-comp2", + output_type=TokenClassifierOutput, + config_class=_CONFIG_FOR_DOC, + expected_loss=0.25, + expected_output=[ + "Lead", + "Lead", + "Lead", + "Position", + "Lead", + "Lead", + "Lead", + "Lead", + "Lead", + "Lead", + "Lead", + "Lead", + ], + ) + # fmt: on + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, + attention_mask: Optional[torch.FloatTensor] = None, + token_type_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + hidden_states = self.dropout(hidden_states) + logits = self.classifier(hidden_states) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + transformer_outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@add_start_docstrings( + """ + The GPT-2 Model transformer with a span classification head on top for extractive question-answering tasks like + SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`). + """, + GPT2_START_DOCSTRING, +) +class GPT2ForQuestionAnswering(GPT2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.transformer = GPT2Model(config) + self.qa_outputs = nn.Linear(config.hidden_size, 2) + + # Model parallel + self.model_parallel = False + self.device_map = None + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING.format("batch_size, sequence_length")) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=QuestionAnsweringModelOutput, + config_class=_CONFIG_FOR_DOC, + real_checkpoint=_CHECKPOINT_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + token_type_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + end_positions: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, QuestionAnsweringModelOutput]: + r""" + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1).to(start_logits.device) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1).to(end_logits.device) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return QuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) \ No newline at end of file diff --git a/ixformer_sdk/train/speedformer/models/llama/__init__.py b/ixformer_sdk/train/speedformer/models/llama/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/models/llama/configuration_llama.py b/ixformer_sdk/train/speedformer/models/llama/configuration_llama.py new file mode 100644 index 00000000..8c44bbd1 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/llama/configuration_llama.py @@ -0,0 +1,191 @@ +# coding=utf-8 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +""" LLaMA model configuration""" + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + + +logger = logging.get_logger(__name__) + +LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP = {} + + +class LlamaConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`LlamaModel`]. It is used to instantiate an LLaMA + model according to the specified arguments, defining the model architecture. Instantiating a configuration with the + defaults will yield a similar configuration to that of the LLaMA-7B. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 32000): + Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`LlamaModel`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 11008): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + num_key_value_heads (`int`, *optional*): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 2048): + The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens, + Llama 2 up to 4096, CodeLlama up to 16384. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + pad_token_id (`int`, *optional*): + Padding token id. + bos_token_id (`int`, *optional*, defaults to 1): + Beginning of stream token id. + eos_token_id (`int`, *optional*, defaults to 2): + End of stream token id. + pretraining_tp (`int`, *optional*, defaults to 1): + Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this + document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is + necessary to ensure exact reproducibility of the pretraining results. Please refer to [this + issue](https://github.com/pytorch/pytorch/issues/76232). + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings + rope_theta (`float`, *optional*, defaults to 10000.0): + The base period of the RoPE embeddings. + rope_scaling (`Dict`, *optional*): + Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling + strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is + `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update + `max_position_embeddings` to the expected new maximum. See the following thread for more information on how + these scaling strategies behave: + https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an + experimental feature, subject to breaking API changes in future versions. + attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): + Whether to use a bias in the query, key, value and output projection layers during self-attention. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + + ```python + >>> from transformers import LlamaModel, LlamaConfig + + >>> # Initializing a LLaMA llama-7b style configuration + >>> configuration = LlamaConfig() + + >>> # Initializing a model from the llama-7b style configuration + >>> model = LlamaModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "llama" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=32000, + hidden_size=4096, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=1, + eos_token_id=2, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self._rope_scaling_validation() + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + def _rope_scaling_validation(self): + """ + Validate the `rope_scaling` configuration. + """ + if self.rope_scaling is None: + return + + if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2: + raise ValueError( + "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, " + f"got {self.rope_scaling}" + ) + rope_scaling_type = self.rope_scaling.get("type", None) + rope_scaling_factor = self.rope_scaling.get("factor", None) + if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]: + raise ValueError( + f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}" + ) + if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0: + raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}") diff --git a/ixformer_sdk/train/speedformer/models/llama/modeling_attn_mask_utils.py b/ixformer_sdk/train/speedformer/models/llama/modeling_attn_mask_utils.py new file mode 100644 index 00000000..67555239 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/llama/modeling_attn_mask_utils.py @@ -0,0 +1,500 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch + + +@dataclass +class AttentionMaskConverter: + """ + A utility attention mask class that allows one to: + - Create a causal 4d mask + - Create a causal 4d mask with slided window + - Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length, + key_value_length) that can be multiplied with attention scores + + Examples: + + ```python + >>> import torch + >>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter + + >>> converter = AttentionMaskConverter(True) + >>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32) + tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]]) + ``` + + Parameters: + is_causal (`bool`): + Whether the attention mask should be a uni-directional (causal) or bi-directional mask. + + sliding_window (`int`, *optional*): + Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer. + """ + + is_causal: bool + sliding_window: int + + def __init__(self, is_causal: bool, sliding_window: Optional[int] = None): + self.is_causal = is_causal + self.sliding_window = sliding_window + + if self.sliding_window is not None and self.sliding_window <= 0: + raise ValueError( + f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`" + ) + + def to_causal_4d( + self, + batch_size: int, + query_length: int, + key_value_length: int, + dtype: torch.dtype, + device: Union[torch.device, "str"] = "cpu", + ) -> Optional[torch.Tensor]: + """ + Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative + bias to upper right hand triangular matrix (causal mask). + """ + if not self.is_causal: + raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.") + + # If shape is not cached, create a new causal mask and cache it + input_shape = (batch_size, query_length) + past_key_values_length = key_value_length - query_length + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if input_shape[-1] > 1 or self.sliding_window is not None: + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + + return causal_4d_mask + + def to_4d( + self, + attention_mask_2d: torch.Tensor, + query_length: int, + dtype: torch.dtype, + key_value_length: Optional[int] = None, + ) -> torch.Tensor: + """ + Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length, + key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is + causal, a causal mask will be added. + """ + input_shape = (attention_mask_2d.shape[0], query_length) + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal: + if key_value_length is None: + raise ValueError( + "This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask." + ) + + past_key_values_length = key_value_length - query_length + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=attention_mask_2d.device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + elif self.sliding_window is not None: + raise NotImplementedError("Sliding window is currently only implemented for causal masking") + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to( + attention_mask_2d.device + ) + + if causal_4d_mask is not None: + expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min) + + # expanded_attn_mask + causal_4d_mask can cause some overflow + expanded_4d_mask = expanded_attn_mask + + return expanded_4d_mask + + @staticmethod + def _make_causal_mask( + input_ids_shape: torch.Size, + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, + ): + """ + Make causal mask used for bi-directional self-attention. + """ + bsz, tgt_len = input_ids_shape + mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device) + mask_cond = torch.arange(mask.size(-1), device=device) + mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0) + + mask = mask.to(dtype) + + if past_key_values_length > 0: + mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1) + + # add lower triangular sliding window mask if necessary + if sliding_window is not None: + diagonal = past_key_values_length - sliding_window + 1 + + context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal) + mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min) + + return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) + + @staticmethod + def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) + + @staticmethod + def _unmask_unattended( + expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float] + ): + # fmt: off + """ + Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when + using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + Details: https://github.com/pytorch/pytorch/issues/110213 + + `expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len]. + `attention_mask` is [bsz, src_seq_len]. + + The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias. + + For example, if `attention_mask` is + ``` + [[0, 0, 1], + [1, 1, 1], + [0, 1, 1]] + ``` + and `expanded_mask` is (e.g. here left-padding case) + ``` + [[[[0, 0, 0], + [0, 0, 0], + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[0, 0, 0], + [0, 1, 0], + [0, 1, 1]]]] + ``` + then the modified `expanded_mask` will be + ``` + [[[[1, 1, 1], <-- modified + [1, 1, 1], <-- modified + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[1, 1, 1], <-- modified + [0, 1, 0], + [0, 1, 1]]]] + ``` + """ + # fmt: on + + # Get the index of the first non-zero value for every sample in the batch. + # In the above example, indices = [[2], [0], [1]]] + tmp = torch.arange(attention_mask.shape[1], 0, -1) + indices = torch.argmax(attention_mask.cpu() * tmp, 1, keepdim=True) + + # Find the batch indexes that have unattended tokens on the leftmost side (e.g. [0, 0, 1, 1, 1]), for which the first rows of the + # expanded mask will be completely unattended. + left_masked_rows = torch.where(indices > 0)[0] + + if left_masked_rows.shape[0] == 0: + return expanded_mask + indices = indices[left_masked_rows] + + max_len = torch.max(indices) + range_tensor = torch.arange(max_len).unsqueeze(0) + range_tensor = range_tensor.repeat(indices.size(0), 1) + + # Avoid unmasking tokens at relevant target positions (on the row axis), by rather unmasking possibly several times the first row that should always be unmasked as we filtered out the batch above. + range_tensor[range_tensor >= indices] = 0 + + # TODO: we may drop support for 3D attention mask as the refactor from Patrick maybe dropped this case + if expanded_mask.dim() == 4: + num_masks = expanded_mask.shape[1] + if num_masks == 1: + # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len] + mask_slice = (left_masked_rows[:, None], 0, range_tensor) + else: + # Broadcast [left_masked_rows, 1, 1], [1, num_masks, 1], [left_masked_rows, 1, max_len] + mask_slice = ( + left_masked_rows[:, None, None], + torch.arange(num_masks)[None, :, None], + range_tensor[:, None, :], + ) + else: + # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len] + mask_slice = (left_masked_rows[:, None], range_tensor) + + expanded_mask[mask_slice] = unmasked_value + + return expanded_mask + + +def _prepare_4d_causal_attention_mask( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + attention_mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + inputs_embeds (`torch.Tensor`): + The embedded inputs as a torch Tensor. + past_key_values_length (`int`): + The length of the key value cache. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + + # 4d mask is passed through the layers + if attention_mask is not None and len(attention_mask.shape) == 2: + attention_mask = attn_mask_converter.to_4d( + attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype + ) + elif attention_mask is not None and len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + else: + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + + return attention_mask + + +# Adapted from _prepare_4d_causal_attention_mask +def _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """ + Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`. + + In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and + `key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks, + allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed). + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + batch_size, query_length = input_shape + + # torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: Fix this as well when using torchdynamo with fullgraph=True. + is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy) + + if attention_mask is not None: + # 4d mask is passed through + if len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype) + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + return attention_mask + + elif not is_tracing and torch.all(attention_mask == 1): + if query_length == 1: + # For query_length == 1, causal attention and bi-directional attention are the same. + attention_mask = None + elif key_value_length == query_length: + attention_mask = None + else: + # Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation + # may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + pass + elif query_length > 1 and key_value_length != query_length: + # See the comment above (https://github.com/pytorch/pytorch/issues/108108). + # Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`. + attention_mask = True + elif is_tracing: + raise ValueError( + 'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.' + ) + + if attention_mask is None: + expanded_4d_mask = None + elif attention_mask is True: + expanded_4d_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + else: + expanded_4d_mask = attn_mask_converter.to_4d( + attention_mask, + input_shape[-1], + dtype=inputs_embeds.dtype, + key_value_length=key_value_length, + ) + + # From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend + # produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213 + # + # This fix is not applied in case we are tracing with torch.jit.trace or symbolic_trace, as _unmask_unattended has a data-dependent + # controlflow that can not be captured properly. + # TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case. + if query_length > 1 and not is_tracing: + expanded_4d_mask = AttentionMaskConverter._unmask_unattended( + expanded_4d_mask, attention_mask, unmasked_value=0.0 + ) + + return expanded_4d_mask + + +def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + batch_size, key_value_length = mask.shape + tgt_len = tgt_len if tgt_len is not None else key_value_length + + # torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: Fix this as well when using torchdynamo with fullgraph=True. + is_tracing = torch.jit.is_tracing() + + if torch.all(mask == 1): + if is_tracing: + pass + elif tgt_len == 1: + # For query_length == 1, causal attention and bi-directional attention are the same. + return None + elif key_value_length == tgt_len: + return None + else: + # Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation + # may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + else: + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _create_4d_causal_attention_mask( + input_shape: Union[torch.Size, Tuple, List], + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, +) -> Optional[torch.Tensor]: + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` + + Args: + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + device (`int`): + The torch device the created mask shall have. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = past_key_values_length + input_shape[-1] + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device + ) + + return attention_mask diff --git a/ixformer_sdk/train/speedformer/models/llama/modeling_llama.py b/ixformer_sdk/train/speedformer/models/llama/modeling_llama.py new file mode 100644 index 00000000..4eca7e09 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/llama/modeling_llama.py @@ -0,0 +1,1415 @@ +# coding=utf-8 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +""" PyTorch LLaMA model.""" +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from .modeling_attn_mask_utils import ( + AttentionMaskConverter, + _prepare_4d_attention_mask, + _prepare_4d_causal_attention_mask, + _prepare_4d_causal_attention_mask_for_sdpa, +) +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13 +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from transformers.utils.import_utils import is_torch_fx_available +from .configuration_llama import LlamaConfig + + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + +# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph. +# It means that the function will not be traced through and simply appear as a node in the graph. +if is_torch_fx_available(): + if not is_torch_greater_or_equal_than_1_13: + import torch.fx + + _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask) + + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "LlamaConfig" + + +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + warnings.warn( + "Calling `transformers.models.llama.modeling_llama._prepare_4d_attention_mask` is deprecated and will be removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask" + ) + return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _make_causal_mask( + input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0 +): + warnings.warn( + "Calling `transformers.models.llama.modeling_llama._make_causal_mask` is deprecated and will be removed in v4.37. Use `transformers.models.llama.modeling_llama.AttentionMaskConverter._make_causal_mask" + ) + return AttentionMaskConverter._make_causal_mask( + input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length + ) + + +class LlamaRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + LlamaRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +ALL_LAYERNORM_LAYERS.append(LlamaRMSNorm) + + +class LlamaRotaryEmbedding(nn.Module): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + super().__init__() + + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # Build here to make `torch.jit.trace` work. + self._set_cos_sin_cache( + seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype() + ) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) + + freqs = torch.outer(t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + def forward(self, x, seq_len=None): + # x: [bs, num_attention_heads, seq_len, head_size] + if seq_len > self.max_seq_len_cached: + self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) + + return ( + self.cos_cached[:seq_len].to(dtype=x.dtype), + self.sin_cached[:seq_len].to(dtype=x.dtype), + ) + + +class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding): + """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev""" + + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): + self.scaling_factor = scaling_factor + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) + t = t / self.scaling_factor + + freqs = torch.outer(t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + +class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding): + """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla""" + + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): + self.scaling_factor = scaling_factor + super().__init__(dim, max_position_embeddings, base, device) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + + if seq_len > self.max_position_embeddings: + base = self.base * ( + (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1) + ) ** (self.dim / (self.dim - 2)) + inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) + + freqs = torch.outer(t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos[position_ids].unsqueeze(unsqueeze_dim) + sin = sin[position_ids].unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class LlamaMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + if self.config.pretraining_tp > 1: + slice = self.intermediate_size // self.config.pretraining_tp + gate_proj_slices = self.gate_proj.weight.split(slice, dim=0) + up_proj_slices = self.up_proj.weight.split(slice, dim=0) + down_proj_slices = self.down_proj.weight.split(slice, dim=1) + + gate_proj = torch.cat( + [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1 + ) + up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1) + + intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2) + down_proj = [ + F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp) + ] + down_proj = sum(down_proj) + else: + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + return down_proj + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class LlamaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.is_causal = True + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) + self._init_rope() + + def _init_rope(self): + if self.config.rope_scaling is None: + self.rotary_emb = LlamaRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + else: + scaling_type = self.config.rope_scaling["type"] + scaling_factor = self.config.rope_scaling["factor"] + if scaling_type == "linear": + self.rotary_emb = LlamaLinearScalingRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + ) + elif scaling_type == "dynamic": + self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + ) + else: + raise ValueError(f"Unknown RoPE scaling type {scaling_type}") + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + bsz, q_len, _ = hidden_states.size() + + if self.config.pretraining_tp > 1: + key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp + query_slices = self.q_proj.weight.split( + (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0 + ) + key_slices = self.k_proj.weight.split(key_value_slicing, dim=0) + value_slices = self.v_proj.weight.split(key_value_slicing, dim=0) + + query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)] + query_states = torch.cat(query_states, dim=-1) + + key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)] + key_states = torch.cat(key_states, dim=-1) + + value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)] + value_states = torch.cat(value_states, dim=-1) + + else: + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + + if self.config.pretraining_tp > 1: + attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2) + o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1) + attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)]) + else: + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class LlamaFlashAttention2(LlamaAttention): + """ + Llama flash attention module. This module inherits from `LlamaAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + # LlamaFlashAttention2 attention does not support output_attentions + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + # overwrite attention_mask with padding_mask + attention_mask = kwargs.pop("padding_mask") + + output_attentions = False + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + dropout_rate = self.attention_dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (LlamaRMSNorm handles it correctly) + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = self._flash_attention_forward( + query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate + ) + + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def _flash_attention_forward( + self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None + ): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`int`, *optional*): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + """ + if not self._flash_attn_uses_top_left_mask: + causal = self.is_causal + else: + # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__. + causal = self.is_causal and query_length != 1 + + # Contains at least one padding token in the sequence + if attention_mask is not None: + batch_size = query_states.shape[0] + query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( + query_states, key_states, value_states, attention_mask, query_length + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) + else: + attn_output = flash_attn_func( + query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal + ) + + return attn_output + + def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = index_first_axis( + key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +class LlamaSdpaAttention(LlamaAttention): + """ + Llama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `LlamaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from LlamaAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if output_attentions: + # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. + logger.warning_once( + "LlamaModel is using LlamaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " + 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attention_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal=self.is_causal and attention_mask is None and q_len > 1, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + return attn_output, None, past_key_value + + +LLAMA_ATTENTION_CLASSES = { + "eager": LlamaAttention, + "flash_attention_2": LlamaFlashAttention2, + "sdpa": LlamaSdpaAttention, +} + + +class LlamaDecoderLayer(nn.Module): + def __init__(self, config: LlamaConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) + + self.mlp = LlamaMLP(config) + self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): + attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, + query_sequence_length, key_sequence_length)` if default attention is used. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + """ + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +LLAMA_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`LlamaConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare LLaMA Model outputting raw hidden-states without any specific head on top.", + LLAMA_START_DOCSTRING, +) +class LlamaPreTrainedModel(PreTrainedModel): + config_class = LlamaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["LlamaDecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +LLAMA_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The bare LLaMA Model outputting raw hidden-states without any specific head on top.", + LLAMA_START_DOCSTRING, +) +class LlamaModel(LlamaPreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] + + Args: + config: LlamaConfig + """ + + def __init__(self, config: LlamaConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._use_sdpa = config._attn_implementation == "sdpa" + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + past_key_values_length = 0 + if use_cache: + use_legacy_cache = not isinstance(past_key_values, Cache) + if use_legacy_cache: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + past_key_values_length = past_key_values.get_usable_length(seq_length) + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if self._use_flash_attention_2: + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + elif self._use_sdpa and not output_attentions: + # output_attentions=True can not be supported when using SDPA, and we fall back on + # the manual implementation that requires a 4D causal mask in all cases. + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + ) + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length + ) + + # embed positions + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class LlamaForCausalLM(LlamaPreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = LlamaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, LlamaForCausalLM + + >>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + if self.config.pretraining_tp > 1: + lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0) + logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)] + logits = torch.cat(logits, dim=-1) + else: + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs + ): + if past_key_values is not None: + if isinstance(past_key_values, Cache): + cache_length = past_key_values.get_seq_length() + past_length = past_key_values.seen_tokens + max_cache_length = past_key_values.get_max_length() + else: + cache_length = past_length = past_key_values[0][0].shape[2] + max_cache_length = None + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as + # input) + if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: + input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + + # If we are about to go beyond the maximum cache length, we need to crop the input attention mask. + if ( + max_cache_length is not None + and attention_mask is not None + and cache_length + input_ids.shape[1] > max_cache_length + ): + attention_mask = attention_mask[:, -max_cache_length:] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + } + ) + return model_inputs + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += ( + tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past), + ) + return reordered_past + + +@add_start_docstrings( + """ + The LLaMa Model transformer with a sequence classification head on top (linear layer). + + [`LlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + LLAMA_START_DOCSTRING, +) +class LlamaForSequenceClassification(LlamaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = LlamaModel(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) diff --git a/ixformer_sdk/train/speedformer/models/qwen2/__init__.py b/ixformer_sdk/train/speedformer/models/qwen2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/models/qwen2/configuration_qwen2.py b/ixformer_sdk/train/speedformer/models/qwen2/configuration_qwen2.py new file mode 100644 index 00000000..b6ca1ed4 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/qwen2/configuration_qwen2.py @@ -0,0 +1,144 @@ +# coding=utf-8 +# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. +# +# 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. +""" Qwen2 model configuration""" + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + + +logger = logging.get_logger(__name__) + +QWEN2_PRETRAINED_CONFIG_ARCHIVE_MAP = { + "Qwen/Qwen2-7B-beta": "https://huggingface.co/Qwen/Qwen2-7B-beta/resolve/main/config.json", +} + + +class Qwen2Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`Qwen2Model`]. It is used to instantiate a + Qwen2 model according to the specified arguments, defining the model architecture. Instantiating a configuration + with the defaults will yield a similar configuration to that of + Qwen2-7B-beta [Qwen/Qwen2-7B-beta](https://huggingface.co/Qwen/Qwen2-7B-beta). + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 151936): + Vocabulary size of the Qwen2 model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`Qwen2Model`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 22016): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer encoder. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer encoder. + num_key_value_heads (`int`, *optional*, defaults to 32): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 32768): + The maximum sequence length that this model might ever be used with. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether the model's input and output word embeddings should be tied. + rope_theta (`float`, *optional*, defaults to 10000.0): + The base period of the RoPE embeddings. + use_sliding_window (`bool`, *optional*, defaults to `False`): + Whether to use sliding window attention. + sliding_window (`int`, *optional*, defaults to 4096): + Sliding window attention (SWA) window size. If not specified, will default to `4096`. + max_window_layers (`int`, *optional*, defaults to 28): + The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + + ```python + >>> from transformers import Qwen2Model, Qwen2Config + + >>> # Initializing a Qwen2 style configuration + >>> configuration = Qwen2Config() + + >>> # Initializing a model from the Qwen2-7B style configuration + >>> model = Qwen2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "qwen2" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=151936, + hidden_size=4096, + intermediate_size=22016, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=32, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_theta=10000.0, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=28, + attention_dropout=0.0, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.use_sliding_window = use_sliding_window + self.sliding_window = sliding_window + self.max_window_layers = max_window_layers + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.attention_dropout = attention_dropout + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/ixformer_sdk/train/speedformer/models/qwen2/modeling_attn_mask_utils.py b/ixformer_sdk/train/speedformer/models/qwen2/modeling_attn_mask_utils.py new file mode 100644 index 00000000..67555239 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/qwen2/modeling_attn_mask_utils.py @@ -0,0 +1,500 @@ +# Copyright 2023 The HuggingFace Team. All rights reserved. +# +# 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. +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch + + +@dataclass +class AttentionMaskConverter: + """ + A utility attention mask class that allows one to: + - Create a causal 4d mask + - Create a causal 4d mask with slided window + - Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length, + key_value_length) that can be multiplied with attention scores + + Examples: + + ```python + >>> import torch + >>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter + + >>> converter = AttentionMaskConverter(True) + >>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32) + tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]]) + ``` + + Parameters: + is_causal (`bool`): + Whether the attention mask should be a uni-directional (causal) or bi-directional mask. + + sliding_window (`int`, *optional*): + Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer. + """ + + is_causal: bool + sliding_window: int + + def __init__(self, is_causal: bool, sliding_window: Optional[int] = None): + self.is_causal = is_causal + self.sliding_window = sliding_window + + if self.sliding_window is not None and self.sliding_window <= 0: + raise ValueError( + f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`" + ) + + def to_causal_4d( + self, + batch_size: int, + query_length: int, + key_value_length: int, + dtype: torch.dtype, + device: Union[torch.device, "str"] = "cpu", + ) -> Optional[torch.Tensor]: + """ + Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative + bias to upper right hand triangular matrix (causal mask). + """ + if not self.is_causal: + raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.") + + # If shape is not cached, create a new causal mask and cache it + input_shape = (batch_size, query_length) + past_key_values_length = key_value_length - query_length + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if input_shape[-1] > 1 or self.sliding_window is not None: + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + + return causal_4d_mask + + def to_4d( + self, + attention_mask_2d: torch.Tensor, + query_length: int, + dtype: torch.dtype, + key_value_length: Optional[int] = None, + ) -> torch.Tensor: + """ + Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length, + key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is + causal, a causal mask will be added. + """ + input_shape = (attention_mask_2d.shape[0], query_length) + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal: + if key_value_length is None: + raise ValueError( + "This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask." + ) + + past_key_values_length = key_value_length - query_length + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=attention_mask_2d.device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + elif self.sliding_window is not None: + raise NotImplementedError("Sliding window is currently only implemented for causal masking") + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to( + attention_mask_2d.device + ) + + if causal_4d_mask is not None: + expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min) + + # expanded_attn_mask + causal_4d_mask can cause some overflow + expanded_4d_mask = expanded_attn_mask + + return expanded_4d_mask + + @staticmethod + def _make_causal_mask( + input_ids_shape: torch.Size, + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, + ): + """ + Make causal mask used for bi-directional self-attention. + """ + bsz, tgt_len = input_ids_shape + mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device) + mask_cond = torch.arange(mask.size(-1), device=device) + mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0) + + mask = mask.to(dtype) + + if past_key_values_length > 0: + mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1) + + # add lower triangular sliding window mask if necessary + if sliding_window is not None: + diagonal = past_key_values_length - sliding_window + 1 + + context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal) + mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min) + + return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) + + @staticmethod + def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) + + @staticmethod + def _unmask_unattended( + expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float] + ): + # fmt: off + """ + Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when + using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + Details: https://github.com/pytorch/pytorch/issues/110213 + + `expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len]. + `attention_mask` is [bsz, src_seq_len]. + + The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias. + + For example, if `attention_mask` is + ``` + [[0, 0, 1], + [1, 1, 1], + [0, 1, 1]] + ``` + and `expanded_mask` is (e.g. here left-padding case) + ``` + [[[[0, 0, 0], + [0, 0, 0], + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[0, 0, 0], + [0, 1, 0], + [0, 1, 1]]]] + ``` + then the modified `expanded_mask` will be + ``` + [[[[1, 1, 1], <-- modified + [1, 1, 1], <-- modified + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[1, 1, 1], <-- modified + [0, 1, 0], + [0, 1, 1]]]] + ``` + """ + # fmt: on + + # Get the index of the first non-zero value for every sample in the batch. + # In the above example, indices = [[2], [0], [1]]] + tmp = torch.arange(attention_mask.shape[1], 0, -1) + indices = torch.argmax(attention_mask.cpu() * tmp, 1, keepdim=True) + + # Find the batch indexes that have unattended tokens on the leftmost side (e.g. [0, 0, 1, 1, 1]), for which the first rows of the + # expanded mask will be completely unattended. + left_masked_rows = torch.where(indices > 0)[0] + + if left_masked_rows.shape[0] == 0: + return expanded_mask + indices = indices[left_masked_rows] + + max_len = torch.max(indices) + range_tensor = torch.arange(max_len).unsqueeze(0) + range_tensor = range_tensor.repeat(indices.size(0), 1) + + # Avoid unmasking tokens at relevant target positions (on the row axis), by rather unmasking possibly several times the first row that should always be unmasked as we filtered out the batch above. + range_tensor[range_tensor >= indices] = 0 + + # TODO: we may drop support for 3D attention mask as the refactor from Patrick maybe dropped this case + if expanded_mask.dim() == 4: + num_masks = expanded_mask.shape[1] + if num_masks == 1: + # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len] + mask_slice = (left_masked_rows[:, None], 0, range_tensor) + else: + # Broadcast [left_masked_rows, 1, 1], [1, num_masks, 1], [left_masked_rows, 1, max_len] + mask_slice = ( + left_masked_rows[:, None, None], + torch.arange(num_masks)[None, :, None], + range_tensor[:, None, :], + ) + else: + # Broadcast [left_masked_rows, 1], [left_masked_rows, max_len] + mask_slice = (left_masked_rows[:, None], range_tensor) + + expanded_mask[mask_slice] = unmasked_value + + return expanded_mask + + +def _prepare_4d_causal_attention_mask( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + attention_mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + inputs_embeds (`torch.Tensor`): + The embedded inputs as a torch Tensor. + past_key_values_length (`int`): + The length of the key value cache. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + + # 4d mask is passed through the layers + if attention_mask is not None and len(attention_mask.shape) == 2: + attention_mask = attn_mask_converter.to_4d( + attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype + ) + elif attention_mask is not None and len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + else: + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + + return attention_mask + + +# Adapted from _prepare_4d_causal_attention_mask +def _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, Tuple, List], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """ + Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`. + + In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and + `key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks, + allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed). + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + batch_size, query_length = input_shape + + # torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: Fix this as well when using torchdynamo with fullgraph=True. + is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy) + + if attention_mask is not None: + # 4d mask is passed through + if len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask.to(inputs_embeds.dtype) + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + return attention_mask + + elif not is_tracing and torch.all(attention_mask == 1): + if query_length == 1: + # For query_length == 1, causal attention and bi-directional attention are the same. + attention_mask = None + elif key_value_length == query_length: + attention_mask = None + else: + # Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore the attention mask, as SDPA causal mask generation + # may be wrong. We will set `is_causal=False` in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + pass + elif query_length > 1 and key_value_length != query_length: + # See the comment above (https://github.com/pytorch/pytorch/issues/108108). + # Ugly: we set it to True here to dispatch in the following controlflow to `to_causal_4d`. + attention_mask = True + elif is_tracing: + raise ValueError( + 'Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model.' + ) + + if attention_mask is None: + expanded_4d_mask = None + elif attention_mask is True: + expanded_4d_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + else: + expanded_4d_mask = attn_mask_converter.to_4d( + attention_mask, + input_shape[-1], + dtype=inputs_embeds.dtype, + key_value_length=key_value_length, + ) + + # From PyTorch 2.1 onwards, F.scaled_dot_product_attention with the memory-efficient attention backend + # produces nans if sequences are completely unattended in the attention mask. Details: https://github.com/pytorch/pytorch/issues/110213 + # + # This fix is not applied in case we are tracing with torch.jit.trace or symbolic_trace, as _unmask_unattended has a data-dependent + # controlflow that can not be captured properly. + # TODO: _unmask_unattended does not work either with torch.compile when using fullgraph=True. We should find a way to detect this case. + if query_length > 1 and not is_tracing: + expanded_4d_mask = AttentionMaskConverter._unmask_unattended( + expanded_4d_mask, attention_mask, unmasked_value=0.0 + ) + + return expanded_4d_mask + + +def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """ + Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)` + + Args: + mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + batch_size, key_value_length = mask.shape + tgt_len = tgt_len if tgt_len is not None else key_value_length + + # torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: Fix this as well when using torchdynamo with fullgraph=True. + is_tracing = torch.jit.is_tracing() + + if torch.all(mask == 1): + if is_tracing: + pass + elif tgt_len == 1: + # For query_length == 1, causal attention and bi-directional attention are the same. + return None + elif key_value_length == tgt_len: + return None + else: + # Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation + # may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + else: + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _create_4d_causal_attention_mask( + input_shape: Union[torch.Size, Tuple, List], + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, +) -> Optional[torch.Tensor]: + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` + + Args: + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + device (`int`): + The torch device the created mask shall have. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = past_key_values_length + input_shape[-1] + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device + ) + + return attention_mask diff --git a/ixformer_sdk/train/speedformer/models/qwen2/modeling_qwen2.py b/ixformer_sdk/train/speedformer/models/qwen2/modeling_qwen2.py new file mode 100644 index 00000000..ec419d26 --- /dev/null +++ b/ixformer_sdk/train/speedformer/models/qwen2/modeling_qwen2.py @@ -0,0 +1,1401 @@ +# coding=utf-8 +# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +""" PyTorch Qwen2 model.""" +import inspect +import math +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from .modeling_attn_mask_utils import _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from .configuration_qwen2 import Qwen2Config + + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters) + + +logger = logging.get_logger(__name__) + + +_CHECKPOINT_FOR_DOC = "Qwen/Qwen2-7B-beta" +_CONFIG_FOR_DOC = "Qwen2Config" + +QWEN2_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "Qwen/Qwen2-7B-beta", + # See all Qwen2 models at https://huggingface.co/models?filter=qwen2 +] + + +# Copied from transformers.models.llama.modeling_llama._get_unpad_data +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Qwen2 +class Qwen2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Qwen2RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Qwen2 +class Qwen2RotaryEmbedding(nn.Module): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + super().__init__() + + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # Build here to make `torch.jit.trace` work. + self._set_cos_sin_cache( + seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype() + ) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) + + freqs = torch.outer(t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) + + def forward(self, x, seq_len=None): + # x: [bs, num_attention_heads, seq_len, head_size] + if seq_len > self.max_seq_len_cached: + self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) + + return ( + self.cos_cached[:seq_len].to(dtype=x.dtype), + self.sin_cached[:seq_len].to(dtype=x.dtype), + ) + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb +def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos[position_ids].unsqueeze(unsqueeze_dim) + sin = sin[position_ids].unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Qwen2 +class Qwen2MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class Qwen2Attention(nn.Module): + """ + Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer + and "Generating Long Sequences with Sparse Transformers". + """ + + def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " + "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.is_causal = True + self.attention_dropout = config.attention_dropout + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.rotary_emb = Qwen2RotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class Qwen2FlashAttention2(Qwen2Attention): + """ + Qwen2 flash attention module, following Qwen2 attention module. This module inherits from `Qwen2Attention` + as the weights of the module stays untouched. The only required change would be on the forward pass + where it needs to correctly call the public API of flash attention and deal with padding tokens + in case the input contains any of them. Additionally, for sliding window attention, we apply SWA only to the bottom + config.max_window_layers layers. + """ + + # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ): + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + # overwrite attention_mask with padding_mask + attention_mask = kwargs.pop("padding_mask") + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1 + cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len) + + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + + use_sliding_windows = ( + _flash_supports_window_size + and getattr(self.config, "sliding_window", None) is not None + and kv_seq_len > self.config.sliding_window + and self.config.use_sliding_window + ) + + if not _flash_supports_window_size: + logger.warning_once( + "The current flash attention version does not support sliding window attention, for a more memory efficient implementation" + " make sure to upgrade flash-attn library." + ) + + if past_key_value is not None: + # Activate slicing cache only if the config has a value `sliding_windows` attribute + cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0 + if ( + getattr(self.config, "sliding_window", None) is not None + and kv_seq_len > self.config.sliding_window + and cache_has_contents + ): + slicing_tokens = 1 - self.config.sliding_window + + past_key = past_key_value[self.layer_idx][0] + past_value = past_key_value[self.layer_idx][1] + + past_key = past_key[:, :, slicing_tokens:, :].contiguous() + past_value = past_value[:, :, slicing_tokens:, :].contiguous() + + if past_key.shape[-2] != self.config.sliding_window - 1: + raise ValueError( + f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got" + f" {past_key.shape}" + ) + + if attention_mask is not None: + attention_mask = attention_mask[:, slicing_tokens:] + attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1) + + cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + dropout_rate = 0.0 if not self.training else self.attention_dropout + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Reashape to the expected shape for Flash Attention + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + attn_output = self._flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + dropout=dropout_rate, + use_sliding_windows=use_sliding_windows, + ) + + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + def _flash_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + query_length, + dropout=0.0, + softmax_scale=None, + use_sliding_windows=False, + ): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`int`, *optional*): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + use_sliding_windows (`bool`, *optional*): + Whether to activate sliding window attention. + """ + if not self._flash_attn_uses_top_left_mask: + causal = self.is_causal + else: + # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__. + causal = self.is_causal and query_length != 1 + + # Decide whether to use SWA or not by layer index. + if use_sliding_windows and self.layer_idx >= self.config.max_window_layers: + use_sliding_windows = False + + # Contains at least one padding token in the sequence + if attention_mask is not None: + batch_size = query_states.shape[0] + query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( + query_states, key_states, value_states, attention_mask, query_length + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + if not use_sliding_windows: + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + else: + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + window_size=(self.config.sliding_window, self.config.sliding_window), + ) + + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) + else: + if not use_sliding_windows: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + else: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=causal, + window_size=(self.config.sliding_window, self.config.sliding_window), + ) + + return attn_output + + # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._upad_input + def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): + batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape + + # On the first iteration we need to properly re-create the padding mask + # by slicing it on the proper place + if kv_seq_len != attention_mask.shape[-1]: + attention_mask_num_tokens = attention_mask.shape[-1] + attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :] + + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + + key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) + value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) + + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +# Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Qwen2 +class Qwen2SdpaAttention(Qwen2Attention): + """ + Qwen2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `Qwen2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from Qwen2Attention.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if output_attentions: + # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. + logger.warning_once( + "Qwen2Model is using Qwen2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " + 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if query_states.device.type == "cuda" and attention_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attention_mask, + dropout_p=self.attention_dropout if self.training else 0.0, + # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1. + is_causal=self.is_causal and attention_mask is None and q_len > 1, + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + return attn_output, None, past_key_value + + +QWEN2_ATTENTION_CLASSES = { + "eager": Qwen2Attention, + "flash_attention_2": Qwen2FlashAttention2, + "sdpa": Qwen2SdpaAttention, +} + + +class Qwen2DecoderLayer(nn.Module): + def __init__(self, config: Qwen2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + if config.use_sliding_window and config._attn_implementation != "flash_attention_2": + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + self.self_attn = QWEN2_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx) + + self.mlp = Qwen2MLP(config) + self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. " + "Please make sure use `attention_mask` instead.`" + ) + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + """ + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +QWEN2_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`Qwen2Config`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare Qwen2 Model outputting raw hidden-states without any specific head on top.", + QWEN2_START_DOCSTRING, +) +class Qwen2PreTrainedModel(PreTrainedModel): + config_class = Qwen2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen2DecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +QWEN2_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +@add_start_docstrings( + "The bare Qwen2 Model outputting raw hidden-states without any specific head on top.", + QWEN2_START_DOCSTRING, +) +class Qwen2Model(Qwen2PreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`] + + Args: + config: Qwen2Config + """ + + def __init__(self, config: Qwen2Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self._attn_implementation = config._attn_implementation + self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + past_key_values_length = 0 + + if use_cache: + use_legacy_cache = not isinstance(past_key_values, Cache) + if use_legacy_cache: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + past_key_values_length = past_key_values.get_usable_length(seq_length) + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if attention_mask is not None and self._attn_implementation == "flash_attention_2" and use_cache: + is_padding_right = attention_mask[:, -1].sum().item() != batch_size + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of Qwen2. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + + if self._attn_implementation == "flash_attention_2": + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + elif self._attn_implementation == "sdpa" and not output_attentions: + # output_attentions=True can not be supported when using SDPA, and we fall back on + # the manual implementation that requires a 4D causal mask in all cases. + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + ) + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + sliding_window=self.config.sliding_window, + ) + + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class Qwen2ForCausalLM(Qwen2PreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = Qwen2Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2ForCausalLM + + >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs + ): + # Omit tokens covered by past_key_values + if past_key_values is not None: + if isinstance(past_key_values, Cache): + cache_length = past_key_values.get_seq_length() + past_length = past_key_values.seen_tokens + max_cache_length = past_key_values.get_max_length() + else: + cache_length = past_length = past_key_values[0][0].shape[2] + max_cache_length = None + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as + # input) + if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: + input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + + # If we are about to go beyond the maximum cache length, we need to crop the input attention mask. + if ( + max_cache_length is not None + and attention_mask is not None + and cache_length + input_ids.shape[1] > max_cache_length + ): + attention_mask = attention_mask[:, -max_cache_length:] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + } + ) + return model_inputs + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + reordered_past = () + for layer_past in past_key_values: + reordered_past += ( + tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past), + ) + return reordered_past + + +@add_start_docstrings( + """ + The Qwen2 Model transformer with a sequence classification head on top (linear layer). + + [`Qwen2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + QWEN2_START_DOCSTRING, +) +class Qwen2ForSequenceClassification(Qwen2PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = Qwen2Model(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(QWEN2_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) diff --git a/ixformer_sdk/train/speedformer/policy/__init__.py b/ixformer_sdk/train/speedformer/policy/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/train/speedformer/policy/baichuan.py b/ixformer_sdk/train/speedformer/policy/baichuan.py new file mode 100644 index 00000000..d24567b7 --- /dev/null +++ b/ixformer_sdk/train/speedformer/policy/baichuan.py @@ -0,0 +1,59 @@ +import warnings +from abc import ABC, abstractmethod +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module + +from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription +from ixformer.train.speedformer.policy.replacer import Replacer + +from ixformer.train.speedformer.models.baichuan.modeling_baichuan import BaichuanModel, DecoderLayer +from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm +from ixformer.train.speedformer.layers.baichuan.attention import BaichuanAttention +from ixformer.train.speedformer.layers.baichuan.mlp import IXFBaichuanMLP + + +class BaichuanReplacer(Replacer): + def __init__(self): + self.policy = {} + + def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={}, + ), + SubModuleReplacementDescription( + suffix="self_attn", + target_module=BaichuanAttention, + kwargs={} + ), + SubModuleReplacementDescription( + suffix="mlp", + target_module=IXFBaichuanMLP, + kwargs={} + ), + ], + target_key="DecoderLayer" + ) + + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="norm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + ], + target_key=BaichuanModel + ) diff --git a/ixformer_sdk/train/speedformer/policy/bloom.py b/ixformer_sdk/train/speedformer/policy/bloom.py new file mode 100644 index 00000000..5380164b --- /dev/null +++ b/ixformer_sdk/train/speedformer/policy/bloom.py @@ -0,0 +1,53 @@ +import warnings +from abc import ABC, abstractmethod +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module + +from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription +from ixformer.train.speedformer.policy.replacer import Replacer + +from ixformer.train.speedformer.models.bloom.modeling_bloom import BloomModel, BloomBlock +from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm +from ixformer.train.speedformer.layers.bloom.attention import BloomFlashAttention + + +class BloomReplacer(Replacer): + def __init__(self): + self.policy = {} + + def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={}, + ), + SubModuleReplacementDescription( + suffix="self_attention", + target_module=BloomFlashAttention, + kwargs={} + ), + ], + target_key="BloomBlock" + ) + + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="ln_f", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + ], + target_key=BloomModel + ) diff --git a/ixformer_sdk/train/speedformer/policy/chatglm.py b/ixformer_sdk/train/speedformer/policy/chatglm.py new file mode 100644 index 00000000..113aeeec --- /dev/null +++ b/ixformer_sdk/train/speedformer/policy/chatglm.py @@ -0,0 +1,57 @@ +from typing import Callable, Dict, List, Union +from torch.nn import Module + +from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription +from ixformer.train.speedformer.policy.replacer import Replacer + +from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm +from ixformer.train.speedformer.layers.chatglm.attention import ChatglmFlashAttention +from ixformer.train.speedformer.layers.chatglm.methods import ChatGLMModel_forward + + +class ChatglmReplacer(Replacer): + def __init__(self): + self.policy = {} + + def module_policy(self) -> Dict[str | Module, List[SubModuleReplacementDescription]]: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="final_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + ], + target_key="GLMTransformer" + ) + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + ], + target_key="GLMBlock" + ) + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="self_attention", + target_module=ChatglmFlashAttention, + kwargs={} + ), + ], + target_key="GLMBlock" + ) + self.append_or_create_method_replacement( + description=[ + {"forward": ChatGLMModel_forward()} + ], + target_key="ChatGLMModel" + ) diff --git a/ixformer_sdk/train/speedformer/policy/gpt2.py b/ixformer_sdk/train/speedformer/policy/gpt2.py new file mode 100644 index 00000000..e450ce8e --- /dev/null +++ b/ixformer_sdk/train/speedformer/policy/gpt2.py @@ -0,0 +1,27 @@ +import torch +import torch.nn as nn +from torch.nn import LayerNorm +from types import ModuleType, MethodType +from abc import ABC + +from ixformer.train.speedformer.models.gpt2.modeling_gpt2 import GPT2FlashAttention2 + +from ixformer.train.speedformer.layers.normalization import replace_layernorm_forward +from ixformer.train.speedformer.layers.gpt2.attention import replace_flash_attn_forward + + +class GPT2Replacer(ABC): + def __init__(self) -> None: + super().__init__() + + @staticmethod + def accelerate(model): + # layer/kernel replace + for name, module in model.named_modules(): + if isinstance(module, LayerNorm): + module.forward = MethodType(replace_layernorm_forward, module) + if isinstance(module, GPT2FlashAttention2): + module._flash_attention_forward = MethodType( + replace_flash_attn_forward, module) + + return model diff --git a/ixformer_sdk/train/speedformer/policy/llama.py b/ixformer_sdk/train/speedformer/policy/llama.py new file mode 100644 index 00000000..08d405b5 --- /dev/null +++ b/ixformer_sdk/train/speedformer/policy/llama.py @@ -0,0 +1,104 @@ +import warnings +import types +from abc import ABC, abstractmethod +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module + +from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription +from ixformer.train.speedformer.policy.replacer import Replacer + +from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm +from ixformer.train.speedformer.layers.llama.attention import LlamaAttention as IXF_LlamaAttention +from ixformer.train.speedformer.layers.llama.mlp import IXFLlamaMLP +from ixformer.train.speedformer.layers.llama.llama_method import LlamaModel_forward, LlamaForCausalLM_forward +from ixformer.train.speedformer.layers.fast_lora.fast_lora import apply_lora_mlp_swiglu + +from peft import PeftType + + +class LlamaReplacer(Replacer): + def __init__(self): + self.policy = {} + + def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={}, + ), + SubModuleReplacementDescription( + suffix="self_attn", + target_module=IXF_LlamaAttention, + kwargs={} + ), + # SubModuleReplacementDescription( + # suffix="mlp", + # target_module=IXFLlamaMLP, + # kwargs={} + # ), + ], + target_key="LlamaDecoderLayer" + ) + + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="norm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + ], + target_key="LlamaModel" + ) + + self.append_or_create_method_replacement( + description=[ + {"forward": LlamaModel_forward()} + ], + target_key="LlamaModel" + ) + self.append_or_create_method_replacement( + description=[ + {"forward": LlamaForCausalLM_forward()} + ], + target_key="LlamaForCausalLM" + ) + + def post_process(self, model: nn.Module): + if model.peft_type != PeftType.LORA: + return + peft_config = model.peft_config + active_adapter = model.active_adapters[0] if \ + hasattr(model, "active_adapters") else model.active_adapter + target_modules = peft_config[active_adapter].target_modules + + # for now, fast_lora only support lora_dropout=0 and bias=None + lora_dropout = model.peft_config[active_adapter].lora_dropout + bias = model.peft_config[active_adapter].bias + + # 首先判断是否可以使用fast_lora + check = lora_dropout == 0 and bias == "none" + + # 其次确定mlp的3个线性层是否在target_modules + mlp_use_fastlora = "gate_proj" in target_modules and "up_proj" in target_modules and "up_proj" in target_modules + + n_mlp = 0 + if check: + if mlp_use_fastlora: + for layer in model.model.model.layers: + layer.mlp.forward = types.MethodType( + apply_lora_mlp_swiglu, layer.mlp) + n_mlp += 1 + + print(f"{len(model.model.model.layers)} layers replace mlp with fast_lora mlp") diff --git a/ixformer_sdk/train/speedformer/policy/qwen2.py b/ixformer_sdk/train/speedformer/policy/qwen2.py new file mode 100644 index 00000000..2b525fbc --- /dev/null +++ b/ixformer_sdk/train/speedformer/policy/qwen2.py @@ -0,0 +1,57 @@ +import warnings +from abc import ABC, abstractmethod +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module + +from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription +from ixformer.train.speedformer.policy.replacer import Replacer +import os +import sys +from ixformer.train.speedformer.layers.normalization import APEXFusedRMSNorm, IXFFusedRMSNorm +from ixformer.train.speedformer.layers.qwen2.attention import QwenAttention as IXF_QwenAttention + + +class Qwen2Replacer(Replacer): + def __init__(self): + self.policy = {} + + def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=APEXFusedRMSNorm, + kwargs={}, + ), + SubModuleReplacementDescription( + suffix="self_attn", + target_module=IXF_QwenAttention, + kwargs={} + ), + ], + target_key="Qwen2DecoderLayer" + ) + + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="norm", + target_module=APEXFusedRMSNorm, + kwargs={} + ), + ], + target_key="Qwen2Model" + ) + + + + \ No newline at end of file diff --git a/ixformer_sdk/train/speedformer/policy/replacer.py b/ixformer_sdk/train/speedformer/policy/replacer.py new file mode 100644 index 00000000..70e3b9c3 --- /dev/null +++ b/ixformer_sdk/train/speedformer/policy/replacer.py @@ -0,0 +1,224 @@ +import warnings +from types import MethodType +from abc import ABC, abstractmethod +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Set, Union +import tabulate + +import torch.nn as nn + +from ixformer.train.speedformer.policy.utils import SubModuleReplacementDescription, ModulePolicyDescription, getattr_, setattr_, print_rank_0 + + +class Replacer(ABC): + def __init__(self): + self.policy = {} + + def module_policy(self) -> Dict[Union[str, nn.Module], List[SubModuleReplacementDescription]]: + r""" + This method returns the module policy, which is a dictionary. The key is the module name or the module object, + and the value is the ModulePolicyDescription object. The ModulePolicyDescription object describes how the module + will be transformed. + """ + + def append_or_create_submodule_replacement( + self, + description: Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]], + target_key: Union[str, nn.Module], + ) -> Dict[Union[str, nn.Module], List]: + r""" + Append or create a new submodule replacement description to the policy for the given key. + + Args: + submodule_replace_desc (Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]]): the submodule replacement description to be appended + policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated + target_key (Union[str, nn.Module]): the key of the policy to be updated + """ + # convert to list + if isinstance(description, SubModuleReplacementDescription): + description = [description] + + # append or create a new description + if target_key in self.policy: + if self.policy[target_key].sub_module_replacement is None: + self.policy[target_key].sub_module_replacement = description + else: + self.policy[target_key].sub_module_replacement.extend( + description) + else: + self.policy[target_key] = ModulePolicyDescription( + sub_module_replacement=description) + + def append_or_create_method_replacement( + self, + description: Dict[str, Callable], + target_key: Union[str, nn.Module], + ) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + r""" + Append or create a new method replacement description to the policy for the given key. + + Args: + description (Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]]): the submodule replacement description to be appended + policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated + target_key (Union[str, nn.Module]): the key of the policy to be updated + """ + if target_key in self.policy: + if self.policy[target_key].method_replacement is None: + self.policy[target_key].method_replacement = description + else: + self.policy[target_key].method_replacement.extend(description) + else: + self.policy[target_key] = ModulePolicyDescription( + method_replacement=description) + + def append_or_create_attribute_replacement( + self, + description: Dict[str, Callable], + target_key: Union[str, nn.Module], + ) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + r""" + Append or create a new method replacement description to the policy for the given key. + + Args: + description (Union[SubModuleReplacementDescription, List[SubModuleReplacementDescription]]): the submodule replacement description to be appended + policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated + target_key (Union[str, nn.Module]): the key of the policy to be updated + """ + if target_key in self.policy: + if self.policy[target_key].attribute_replacement is None: + self.policy[target_key].attribute_replacement = description + else: + self.policy[target_key].attribute_replacement.extend( + description) + else: + self.policy[target_key] = ModulePolicyDescription( + attribute_replacement=description) + + def accelerate(self, model) -> None: + r""" + Replace the module according to the policy, and replace the module one by one + + Args: + model (:class:`torch.nn.Module`): The model to shard + """ + self.module_policy() + self.module_replace = [] + for layer_cls, module_description in self.policy.items(): + self.replace_sub_module( + model, layer_cls, module_description.sub_module_replacement) + self._replace_method( + model, layer_cls, module_description.method_replacement) + print_rank_0(tabulate.tabulate(self.module_replace, headers=[ + "old_layer", "new_layer"], tablefmt="psql")) + return model + + def replace_sub_module( + self, + module: nn.Module, + origin_cls: Union[str, nn.Module], + sub_module_replacement: List[SubModuleReplacementDescription], + ) -> None: + r""" + Reverse the replace layer operation + """ + if not sub_module_replacement: + return + + if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or ( + module.__class__ == origin_cls + ): + for description in sub_module_replacement: + suffix = description.suffix + target_module = description.target_module + kwargs = {} if description.kwargs is None else description.kwargs + + assert target_module is not None, "target_module should not be None" + + native_sub_module = getattr_(module, suffix, ignore=True) + + assert not isinstance( + native_sub_module, target_module + ), f"The module with suffix {suffix} has been replaced, please check the policy" + + # if it is None and we are allowed to ignore this module + # just skip + if description.ignore_if_not_exist and native_sub_module is None: + continue + try: + replace_layer = target_module.from_native_module( + native_sub_module, **kwargs) + except Exception as e: + raise RuntimeError( + f"Failed to replace {suffix} of type {native_sub_module.__class__.__qualname__}" + f" with {target_module.__qualname__} with the exception: {e}. " + "Please check your model configuration or sharding policy, you can set up an issue for us to help you as well." + ) + + setattr_(module, suffix, replace_layer) + self.module_replace.append( + [native_sub_module.__class__.__qualname__, target_module.__qualname__]) + + for name, child in module.named_children(): + self.replace_sub_module( + child, + origin_cls, + sub_module_replacement, + ) + + def _replace_method(self, module: nn.Module, origin_cls: Union[str, nn.Module], method_replacement: List[Dict[str, Callable]]): + if not method_replacement: + return + + if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or ( + module.__class__ == origin_cls + ): + for method in method_replacement: + for method_name, new_method in method.items(): + # bind the new method to the module + bound_method = MethodType(new_method, module) + setattr(module, method_name, bound_method) + + for name, child in module.named_children(): + self._replace_method( + child, + origin_cls, + method_replacement, + ) + + def _replace_attr( + self, + module: nn.Module, + origin_cls: Union[str, nn.Module], + attr_replacement: List[Dict[str, Any]], + ) -> None: + r""" + Replace the attribute of the layer + + Args: + module (:class:`torch.nn.Module`): The object of layer to shard + attr_replacement (Dict): The attribute dict to modify + """ + if not attr_replacement: + return + + if (isinstance(origin_cls, str) and origin_cls == module.__class__.__name__) or ( + module.__class__ == origin_cls + ): + for attr in attr_replacement: + for module_attr, target_attr in attr.items(): + native_attr = getattr_(module, module_attr, ignore=False) + if isinstance(native_attr, type): + replace_attr = target_attr.from_native_attr( + native_attr) + setattr_(module, module_attr, + replace_attr, ignore=False) + else: + setattr_(module, module_attr, + target_attr, ignore=False) + + for name, child in module.named_children(): + self._replace_attr( + child, + origin_cls, + attr_replacement, + ) diff --git a/ixformer_sdk/train/speedformer/policy/utils.py b/ixformer_sdk/train/speedformer/policy/utils.py new file mode 100644 index 00000000..f377145c --- /dev/null +++ b/ixformer_sdk/train/speedformer/policy/utils.py @@ -0,0 +1,156 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Union +import re +import torch +import torch.nn as nn + + +@dataclass +class SubModuleReplacementDescription: + r""" + Describe how a submodule will be replaced + + Args: + suffix (str): used to get the submodule object + target_module (ParallelModule): specifies the module class used to replace to submodule + kwargs (Dict[str, Any]): the dictionary used to pass extra arguments to the `ParallelModule.from_native_module` method. + ignore_if_not_exist (bool): if the submodule does not exist, ignore it or raise an exception + """ + + suffix: str + target_module: nn.Module + kwargs: Dict[str, Any] = None + ignore_if_not_exist: bool = False + + +@dataclass +class ModulePolicyDescription: + "copy from colossalai, for now sub_module_replacement and method_replacement is used" + r""" + Describe how the attributes and parameters will be transformed in a policy. + + Args: + attribute_replacement (Dict[str, Any]): key is the attribute name, value is the attribute value after sharding + param_replacement (List[Callable]): a list of functions to perform in-place param replacement. The function + must receive only one arguments: module. One example is + + ```python + def example_replace_weight(module: torch.nn.Module): + weight = module.weight + new_weight = shard_rowwise(weight, process_group) + module.weight = torch.nn.Parameter(new_weight) + ``` + sub_module_replacement (List[SubModuleReplacementDescription]): each element in the list is a SubModuleReplacementDescription + object which specifies the module to be replaced and the target module used to replacement. + method_replace (Dict[str, Callable]): key is the method name, value is the method for replacement + """ + + attribute_replacement: List[Dict[str, Any]] = None + param_replacement: List[Callable] = None + sub_module_replacement: List[SubModuleReplacementDescription] = None + method_replacement: List[Dict[str, Callable]] = None + + +def getattr_(obj, attr: str, ignore: bool = False): + r""" + Get the object's multi sublevel attr + + Args: + obj (object): The object to set + attr (str): The multi level attr to set + ignore (bool): Whether to ignore when the attr doesn't exist + """ + + attrs = attr.split(".") + for a in attrs: + try: + obj = get_obj_list_element(obj, a) + except AttributeError: + if ignore: + return None + raise AttributeError( + f"Object {obj.__class__.__name__} has no attribute {attr}") + return obj + + +def get_obj_list_element(obj, attr: str): + r""" + Get the element of the list in the object + + If the attr is a normal attribute, return the attribute of the object. + If the attr is a index type, return the element of the index in the list, like `layers[0]`. + + Args: + obj (Object): The object to get + attr (str): The suffix of the attribute to get + + """ + re_pattern = r"\[\d+\]" + prog = re.compile(re_pattern) + result = prog.search(attr) + if result: + matched_brackets = result.group() + matched_index = matched_brackets.replace("[", "") + matched_index = matched_index.replace("]", "") + attr_ = attr.replace(matched_brackets, "") + container_obj = getattr(obj, attr_) + obj = container_obj[int(matched_index)] + else: + obj = getattr(obj, attr) + return obj + + +def setattr_(obj, attr: str, value, ignore: bool = False): + r""" + Set the object's multi sublevel attr to value, if ignore, ignore when it doesn't exist + + Args: + obj (object): The object to set + attr (str): The multi level attr to set + value (Any): The value to set + ignore (bool): Whether to ignore when the attr doesn't exist + """ + + attrs = attr.split(".") + for a in attrs[:-1]: + try: + obj = get_obj_list_element(obj, a) + except AttributeError: + if ignore: + return + raise AttributeError( + f"Object {obj.__class__.__name__} has no attribute {attr}") + set_obj_list_element(obj, attrs[-1], value) + + +def set_obj_list_element(obj, attr: str, value): + r""" + Set the element to value of a list object + + It used like set_obj_list_element(obj, 'layers[0]', new_layer), it will set obj.layers[0] to value + + Args: + obj (object): The object to set + attr (str): the string including a list index like `layers[0]` + """ + re_pattern = r"\[\d+\]" + prog = re.compile(re_pattern) + result = prog.search(attr) + if result: + matched_brackets = result.group() + matched_index = matched_brackets.replace("[", "") + matched_index = matched_index.replace("]", "") + attr_ = attr.replace(matched_brackets, "") + container_obj = getattr(obj, attr_) + container_obj[int(matched_index)] = value + else: + setattr(obj, attr, value) + + +def print_rank_0(message): + if torch.distributed.is_initialized(): + if torch.distributed.get_rank() == 0: + print(message, flush=True) + else: + print(message, flush=True) diff --git a/ixformer_sdk/train/speedformer/speedformer.py b/ixformer_sdk/train/speedformer/speedformer.py new file mode 100644 index 00000000..cb578241 --- /dev/null +++ b/ixformer_sdk/train/speedformer/speedformer.py @@ -0,0 +1,25 @@ +import torch +import torch.nn as nn +from abc import ABC +from ixformer.train.speedformer.model_replacer_mapping import ModelMapping + +# 外部接口 +class SpeedFormer(ABC): + def __init__(self) -> None: + super().__init__() + self.replacer = None + + + def accelerate(self, model): + if model.config.model_type in ModelMapping: + self.replacer = ModelMapping[model.config.model_type]() + accelerate_model = self.replacer.accelerate(model) + else: + Warning(f"Warning: model '{model.config.model_type}' is not supported now.") + accelerate_model = model + + return accelerate_model + + + def post_process(self, model): + self.replacer.post_process(model) diff --git a/ixformer_sdk/utils/__init__.py b/ixformer_sdk/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ixformer_sdk/utils/benchmark/__init__.py b/ixformer_sdk/utils/benchmark/__init__.py new file mode 100644 index 00000000..730ede6c --- /dev/null +++ b/ixformer_sdk/utils/benchmark/__init__.py @@ -0,0 +1 @@ +from .timer import Benchmark, BenchmarkTimer diff --git a/ixformer_sdk/utils/benchmark/cuda_benchmark.py b/ixformer_sdk/utils/benchmark/cuda_benchmark.py new file mode 100644 index 00000000..995b2bcd --- /dev/null +++ b/ixformer_sdk/utils/benchmark/cuda_benchmark.py @@ -0,0 +1,69 @@ +import time +from collections import namedtuple, OrderedDict +from typing import List, Dict, Any + +import tabulate +import torch +import torch.distributed as dist + +DeviceTime = namedtuple("DeviceTime", ["cpu", "gpu"]) + + +class Functor: + + def __init__(self, fn, *args, **kwargs): + self.fn = fn + self.args = args + self.kwargs = kwargs + + def __call__(self): + return self.fn(*self.args, **self.kwargs) + + +def cuda_timeit(fn: Functor, dist_barrier=False) -> DeviceTime: + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + stop = torch.cuda.Event(enable_timing=True) + start.record(torch.cuda.current_stream()) + + t0 = time.time() + fn() + t1 = time.time() + + stop.record(torch.cuda.current_stream()) + + torch.cuda.synchronize() + if dist_barrier: + dist.barrier() + + gpu_time = start.elapsed_time(stop) + cpu_time = t1 - t0 + return DeviceTime(cpu_time, gpu_time) + + +def cuda_benchmark(fn: Functor, num_repeated=10, num_warmup=1, dist_barrier=False) -> DeviceTime: + [fn() for _ in range(num_warmup)] + times = [cuda_timeit(fn, dist_barrier=dist_barrier) for _ in range(num_repeated)] + times.sort(key=lambda t: t.gpu) + + if num_repeated >= 10: + times = times[3:-3] + + avg_gpu_time = sum([t.gpu for t in times]) / len(times) + avg_cpu_time = sum([t.cpu for t in times]) / len(times) + + return DeviceTime(avg_cpu_time * 1000, avg_gpu_time) + + +def show_benchmark_results(times: List[DeviceTime], extra_info: Dict[Any, List]=None): + data = extra_info or OrderedDict() + + if len(times) != 0: + cpu_times = [round(t.cpu, 6) for t in times] + gpu_times = [round(t.gpu, 6) for t in times] + + data["CPU Time(ms)"] = cpu_times + data["GPU Time(ms)"] = gpu_times + + print(tabulate.tabulate(extra_info, headers=data.keys())) diff --git a/ixformer_sdk/utils/benchmark/timer.py b/ixformer_sdk/utils/benchmark/timer.py new file mode 100644 index 00000000..3878d66f --- /dev/null +++ b/ixformer_sdk/utils/benchmark/timer.py @@ -0,0 +1,130 @@ +import time +from collections import OrderedDict +from typing import Callable + +from tabulate import tabulate +from tqdm import tqdm +import torch + + +class BenchmarkTimer: + def __init__(self): + self.reset() + + def reset(self): + self.start_time = None + self.end_time = None + self.running_times = [] + + def __enter__(self): + self.start_time = time.perf_counter() + + def __exit__(self, exc_type, exc_val, exc_tb): + self.end_time = time.perf_counter() + self.running_times.append(self.end_time - self.start_time) + + +class Benchmark: + def __init__( + self, + warmup: int = None, + number: int = 100, + timer=None, + description: str = None, + show_progress: bool = False, + fn_desc_key: str = "fn_desc", + sync: bool = True, + ): + if warmup is None: + warmup = int(number // 100) + 10 + self.warmup = warmup + self.number = number + self.description = description + self.show_progress = show_progress + self.fn_desc_key = fn_desc_key + self.sync = sync + + if timer is None: + timer = BenchmarkTimer() + self.timer = timer + + self.reset() + + def reset(self): + self.results = OrderedDict() + self._run_index = 0 + self._fn_name = None + + def run(self, fn, *args, **kwargs): + self._run_index += 1 + + if self.fn_desc_key in kwargs: + self.set_fn_name(kwargs[self.fn_desc_key]) + kwargs.pop(self.fn_desc_key) + key = self._get_fn_key(fn) + + # warmup + self._run_fn(False, fn, *args, **kwargs) + + # get running times + results = self._run_fn(True, fn, *args, **kwargs) + self.results[key] = results + + return results + + def set_fn_name(self, name): + self._fn_name = name + + def _run_fn(self, benchmark: bool, fn: Callable, *args, **kwargs): + self.timer.reset() + n = self.number if benchmark else self.warmup + if self.show_progress and benchmark: + progress = tqdm(range(n), desc=self._get_fn_key(fn)) + else: + progress = range(n) + + torch.cuda.synchronize() + + for _ in progress: + with self.timer: + fn(*args, **kwargs) + + if self.sync: + torch.cuda.synchronize() + + return self.timer.running_times + + def _get_fn_key(self, fn: Callable): + if self._fn_name is not None: + return self._fn_name + + if hasattr(fn, "__name__"): + fn_name = fn.__name__ + else: + fn_name = str(fn) + + return f"{fn_name}_{self._run_index}" + + def render(self) -> str: + head = [""] + list(self.results.keys()) + total = ["Total (s)"] + [sum(times) for times in self.results.values()] + mean = ["Mean (s)"] + [_t / self.number for _t in total[1:]] + min_ = ["Min (s)"] + [min(times) for times in self.results.values()] + max_ = ["Max (s)"] + [max(times) for times in self.results.values()] + count = ["Count"] + [len(list(times)) for times in self.results.values()] + + return tabulate( + headers=head, + tabular_data=[total, mean, min_, max_, count], + numalign="right", + ) + + def print_caption(self): + if self.description is not None: + caption = "\n" + "=" * 60 + "\n" + caption += f"= {self.description}" + "\n" + caption += "=" * 60 + "\n" + print(caption) + + def print(self): + print(self.render()) diff --git a/ixformer_sdk/utils/object.py b/ixformer_sdk/utils/object.py new file mode 100644 index 00000000..52dea523 --- /dev/null +++ b/ixformer_sdk/utils/object.py @@ -0,0 +1,227 @@ +import inspect +from typing import Any, Callable, Dict, Mapping, Union + +__all__ = [ + "isfunction", + "iscallable", + "get_obj_name", + "isimmutable_var", + "get_self_from", + "get_obj_funcs", + "recurse_getattr", + "recurse_find_by_key", + "set_value_by_cascasde_key", + "flatten_container", + "flatten_dict", + "get_func_argspec", + "get_obj_attr", + "get_namedtuple_fields", + "get_namedtuple_defaults", + "isnamedtuple", + "namedtype_to_dict", +] + + +def isfunction(f): + return ( + inspect.isfunction(f) or inspect.ismethod(f) or inspect.isbuiltin(f) + ) and not inspect.isclass(f) + + +def iscallable(fn) -> bool: + return any( + [ + callable(fn), + inspect.isfunction(fn), + inspect.ismethod(fn), + inspect.isbuiltin(fn), + ] + ) + + +def get_obj_name(obj, containe_module=False): + mod_name = None + if inspect.isclass(obj): + obj_name = obj.__name__ + if hasattr(obj, "__module__") and containe_module: + mod_name = obj.__module__ + elif hasattr(obj, "__name__"): + obj_name = obj.__name__ + elif hasattr(obj, "__class__"): + obj_name = obj.__class__.__name__ + if hasattr(obj.__class__, "__module__") and containe_module: + mod_name = obj.__class__.__module__ + else: + obj_name = str(obj) + + if containe_module and mod_name is None: + if hasattr(obj, "__module__"): + mod_name = obj.__module__ + + if mod_name is None: + return obj_name + else: + return f"{mod_name}.{obj_name}" + + +def isimmutable_var(var): + if var is None: + return True + + if inspect.isclass(var): + var_cls = var + else: + var_cls = type(var) + + return var_cls in [int, float, tuple, str, None] + + +def get_self_from(obj): + if hasattr(obj, "__self__"): + return obj.__self__ + raise AttributeError(f"Not found attribute `self` in {obj}.") + + +def get_obj_funcs(obj) -> Dict[str, Callable]: + attrs = dir(obj) + funcs = dict() + for attr in attrs: + fn = getattr(obj, attr) + if iscallable(fn): + funcs[attr] = fn + + return funcs + + +def recurse_find_by_key(container: dict, key: Union[str, list], default=None): + if isinstance(key, str): + key = key.split(".") + + if not isinstance(key, (tuple, list)): + raise RuntimeError(f"Please give the type str or list, but get ({type(key)}).") + + value = default + _cnt = container + for k in key: + if k not in _cnt: + return default + value = _cnt[k] + _cnt = value + if _cnt is None: + return default + + return value + + +def set_value_by_cascasde_key(container: dict, key: str, value: Any): + if isinstance(key, str): + key = key.split(".") + + if not isinstance(key, (tuple, list)): + raise RuntimeError(f"Please give the type str or list, but get ({type(key)}).") + + _cnt = container + for k in key[:-1]: + if k not in _cnt: + _cnt[k] = dict() + _cnt = _cnt[k] + _cnt[key[-1]] = value + return container + + +def flatten_dict(d: dict, preffix="", out=None): + if out is None: + out = dict() + for k, v in d.items(): + if isinstance(v, Mapping): + flatten_dict(v, f"{preffix}{k}.", out) + else: + out[preffix + k] = v + + return out + + +def flatten_container(container: Union[list, dict]): + outs = [] + + def _flatten_list(cnt: list): + for item in cnt: + if isinstance(item, (tuple, list)): + _flatten_list(item) + elif isinstance(item, Mapping): + _flatten_dict(item) + else: + outs.append(item) + + def _flatten_dict(cnt: Dict): + for key, item in cnt.items(): + if isinstance(item, (tuple, list)): + _flatten_list(item) + elif isinstance(item, Mapping): + _flatten_dict(item) + else: + outs.append(item) + + if isinstance(container, (tuple, list)): + _flatten_list(container) + elif isinstance(container, dict): + _flatten_dict(container) + else: + outs.append(container) + + return outs + + +def get_func_argspec(func) -> inspect.FullArgSpec: + return inspect.getfullargspec(func) + + +def get_obj_attr(obj, attr, default=None): + if isinstance(obj, Mapping): + return obj.get(attr, default) + return getattr(obj, attr, default) + + +def isnamedtuple(obj): + if not inspect.isclass(obj) or not issubclass(obj, tuple): + return False + + if hasattr(obj, "_fields") and hasattr(obj, "_replace"): + if ( + hasattr(obj._replace, "__module__") + and obj._replace.__module__ == "collections" + ): + return True + + return False + + +def get_namedtuple_fields(t): + if not inspect.isclass(t): + t = type(t) + + if not isnamedtuple(t): + raise RuntimeError(f"{t} is not a namedtuple object") + + return t._fields + + +def get_namedtuple_defaults(t) -> dict: + return t._field_defaults + + +def namedtype_to_dict(t): + return t._asdict() + + +def recurse_getattr(obj, attr: str, sep="."): + attrs = attr.split(sep) + idx = 0 + cur_obj = obj + while idx < len(attrs): + cur_obj = getattr(cur_obj, attrs[idx]) + idx += 1 + + if cur_obj == obj: + return None + return cur_obj diff --git a/ixformer_sdk/utils/seed.py b/ixformer_sdk/utils/seed.py new file mode 100644 index 00000000..dff01700 --- /dev/null +++ b/ixformer_sdk/utils/seed.py @@ -0,0 +1,16 @@ +import random + +import numpy as np + + +def manual_seed(seed=41): + random.seed(seed) + np.random.seed(seed) + try: + import torch + + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + except: + pass diff --git a/ixformer_sdk/version.py b/ixformer_sdk/version.py new file mode 100644 index 00000000..32156fda --- /dev/null +++ b/ixformer_sdk/version.py @@ -0,0 +1,4 @@ +__version__ = "0.6.0" +torch = "2.4.1+corex.4.3.8" +git_commit = "710dc9444aa2" +vllm = "None" diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/CMakeLists.txt b/upstream_ref/xllm_latest/core/layers/npu_torch/CMakeLists.txt new file mode 100755 index 00000000..83b57c02 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/CMakeLists.txt @@ -0,0 +1,28 @@ +include(cc_library) + +cc_library( + NAME + npu_torch_layers + HDRS + fused_moe.h + attention.h + qwen3_gated_delta_net_base.h + qwen3_next_attention.h + qwen3_next_gated_delta_net.h + qwen3_5_gated_delta_net.h + qwen3_next_hybrid_decoder_layer_base.h + qwen3_next_decoder_layer_impl.h + qwen3_5_decoder_layer_impl.h + SRCS + fused_moe.cpp + attention.cpp + qwen3_gated_delta_net_base.cpp + qwen3_next_attention.cpp + qwen3_next_gated_delta_net.cpp + qwen3_next_hybrid_decoder_layer_base.cpp + qwen3_5_gated_delta_net.cpp + qwen3_next_decoder_layer_impl.cpp + qwen3_5_decoder_layer_impl.cpp + DEPS + :common_layers +) diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/attention.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/attention.cpp new file mode 100644 index 00000000..eb2b7c6c --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/attention.cpp @@ -0,0 +1,152 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "attention.h" + +#include "kernels/npu/npu_ops_api.h" +#include "kernels/ops_api.h" + +DECLARE_bool(enable_chunked_prefill); +namespace xllm { +namespace layer { + +AttentionImpl::AttentionImpl(int64_t num_heads, + int64_t head_size, + float scale, + int64_t num_kv_heads, + int64_t sliding_window) + : num_heads_(num_heads), + head_size_(head_size), + num_kv_heads_(num_kv_heads), + sliding_window_(sliding_window), + scale_(scale) { + if (sliding_window_ > -1) { + sliding_window_ = sliding_window_ - 1; + } +} + +std::tuple> AttentionImpl::forward( + const AttentionMetadata& attn_metadata, + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + KVCache& kv_cache) { + std::optional output_lse = std::nullopt; + torch::Tensor output = torch::empty_like(query); + + if (attn_metadata.is_dummy) { + return std::make_tuple(output, output_lse); + } + + bool only_prefill = + attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; + + torch::Tensor k_cache = kv_cache.get_k_cache(); + torch::Tensor v = value.view({-1, num_kv_heads_, head_size_}); + std::optional v_cache = kv_cache.get_v_cache(); + + // Reshape and cache key/value + xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params; + reshape_paged_cache_params.key = key.view({-1, num_kv_heads_, head_size_}); + reshape_paged_cache_params.value = v; + reshape_paged_cache_params.k_cache = k_cache; + reshape_paged_cache_params.v_cache = v_cache; + reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping; + xllm::kernel::reshape_paged_cache(reshape_paged_cache_params); + + if (only_prefill) { + prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata); + } else { + decoder_forward(query, output, k_cache, v_cache, attn_metadata); + } + + output = output.view({-1, num_heads_ * head_size_}); + return {output, output_lse}; +} + +void AttentionImpl::prefill_forward(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + query = query.view({-1, num_heads_, head_size_}); + output = output.view({-1, num_heads_, head_size_}); + + if (attn_metadata.is_prefill) { + key = key.view({-1, num_kv_heads_, head_size_}); + value = value.view({-1, num_kv_heads_, head_size_}); + + xllm::kernel::npu::batch_prefill(query, + key, + value, + attn_metadata.attn_mask, + attn_metadata.kv_seq_lens_host, + scale_, + output); + } else if (attn_metadata.is_chunked_prefill) { + xllm::kernel::npu::batch_prefill(query, + k_cache, + v_cache.value(), + attn_metadata.attn_mask, + attn_metadata.kv_seq_lens_host, + scale_, + output); + } +} + +void AttentionImpl::decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + query = query.view({-1, 1, num_heads_, head_size_}); + output = output.view({-1, 1, num_heads_, head_size_}); + + torch::Tensor kv_seq_lens; + if (attn_metadata.kv_seq_lens_host.defined()) { + kv_seq_lens = attn_metadata.kv_seq_lens_host; + } else { + // Fallback if host tensor isn't prepared. + kv_seq_lens = attn_metadata.kv_seq_lens; + } + + if (attn_metadata.paged_attention_tiling_data.defined()) { + // Use CustomPagedAttention for ACL graph mode to avoid .to(kCPU) operations + + xllm::kernel::npu::batch_decode_acl_graph( + query, + k_cache, + v_cache.value_or(torch::Tensor()), + scale_, + attn_metadata.block_table, + kv_seq_lens, + attn_metadata.paged_attention_tiling_data, + output); + } else { + // Standard PagedAttention path + xllm::kernel::npu::batch_decode(query, + k_cache, + v_cache.value_or(torch::Tensor()), + scale_, + attn_metadata.block_table, + kv_seq_lens, + output); + } +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/attention.h b/upstream_ref/xllm_latest/core/layers/npu_torch/attention.h new file mode 100644 index 00000000..f3a9c0e1 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/attention.h @@ -0,0 +1,70 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include + +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_input_params.h" +#include "layers/common/attention_metadata.h" + +namespace xllm { +namespace layer { + +class AttentionImpl : public torch::nn::Module { + public: + AttentionImpl() = default; + + AttentionImpl(int64_t num_heads, + int64_t head_size, + float scale, + int64_t num_kv_heads, + int64_t sliding_window); + + std::tuple> forward( + const AttentionMetadata& attn_metadata, + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + KVCache& kv_cache); + + void prefill_forward(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + + void decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + + private: + int64_t num_heads_; + int64_t head_size_; + float scale_; + int64_t num_kv_heads_; + int64_t sliding_window_; +}; +TORCH_MODULE(Attention); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.cpp new file mode 100644 index 00000000..b13d6d6f --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.cpp @@ -0,0 +1,513 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "fused_moe.h" + +#include + +#include +#include + +#include "framework/parallel_state/parallel_state.h" +#include "kernels/ops_api.h" + +namespace xllm { +namespace layer { + +namespace { +// Generic local tensor helpers. +torch::Tensor create_group_gemm_output( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& group_list, + torch::ScalarType dtype = torch::ScalarType::BFloat16) { + torch::TensorOptions target_options = a.options().dtype(dtype); + if (b.dim() != 2) { + return torch::empty({a.size(0), b.size(1)}, target_options); + } + return torch::empty({group_list.size(0), a.size(0), b.size(0)}, + target_options); +} + +torch::Tensor get_tensor_with_weight_suffix(const StateDict& state_dict, + const std::string& tensor_name) { + auto tensor = state_dict.get_tensor(tensor_name); + if (!tensor.defined()) { + tensor = state_dict.get_tensor(tensor_name + ".weight"); + } + return tensor; +} + +torch::Tensor slice_expert_weights(const torch::Tensor& weight, + int64_t start_expert_id, + int64_t num_experts_per_rank) { + return weight + .slice(0, start_expert_id, start_expert_id + num_experts_per_rank) + .contiguous(); +} + +// Qwen3.5-MoE fused checkpoint fallback helpers. +bool load_fused_gate_up_fallback(const StateDict& state_dict, + int64_t rank, + int64_t world_size, + int64_t start_expert_id, + int64_t num_experts_per_rank, + torch::Tensor& w13) { + auto fused_gate_up = + get_tensor_with_weight_suffix(state_dict, "gate_up_proj"); + if (!fused_gate_up.defined()) { + return false; + } + + if (world_size > 1) { + CHECK_EQ(fused_gate_up.size(1) % 2, 0) + << "gate_up_proj dim1 must be even, got " << fused_gate_up.size(1); + const int64_t full_intermediate = fused_gate_up.size(1) / 2; + CHECK_EQ(full_intermediate % world_size, 0) + << "gate_up_proj intermediate dim is not divisible by world_size"; + const int64_t inter_shard = full_intermediate / world_size; + + auto gate_full = fused_gate_up.slice(1, 0, full_intermediate); + auto up_full = + fused_gate_up.slice(1, full_intermediate, full_intermediate * 2); + auto gate_shard = + gate_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard); + auto up_shard = + up_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard); + fused_gate_up = torch::cat({gate_shard, up_shard}, 1); + } + + auto gate_up_slice = slice_expert_weights( + fused_gate_up, start_expert_id, num_experts_per_rank); + CHECK_EQ(w13.sizes(), gate_up_slice.sizes()) + << "weight size mismatch for " << state_dict.prefix() + << "experts.gate_up_proj"; + w13.copy_(gate_up_slice); + return true; +} + +bool load_fused_down_fallback(const StateDict& state_dict, + int64_t rank, + int64_t world_size, + int64_t start_expert_id, + int64_t num_experts_per_rank, + torch::Tensor& w2) { + auto fused_down = get_tensor_with_weight_suffix(state_dict, "down_proj"); + if (!fused_down.defined()) { + return false; + } + + if (world_size > 1) { + CHECK_EQ(fused_down.size(2) % world_size, 0) + << "down_proj dim2 is not divisible by world_size"; + const int64_t down_shard = fused_down.size(2) / world_size; + fused_down = + fused_down.slice(2, rank * down_shard, (rank + 1) * down_shard); + } + + auto down_slice = + slice_expert_weights(fused_down, start_expert_id, num_experts_per_rank); + CHECK_EQ(w2.sizes(), down_slice.sizes()) + << "weight size mismatch for " << state_dict.prefix() + << "experts.down_proj"; + w2.copy_(down_slice); + return true; +} + +} // namespace + +FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options) + : num_total_experts_(model_args.n_routed_experts()), + topk_(model_args.num_experts_per_tok()), + hidden_size_(model_args.hidden_size()), + n_shared_experts_(model_args.n_shared_experts()), + is_gated_(moe_args.is_gated), + renormalize_(model_args.norm_topk_prob() ? 1 : 0), + hidden_act_(model_args.hidden_act()), + is_smoothquant_(false), + quant_args_(quant_args), + parallel_args_(parallel_args), + options_(options), + tp_pg_(parallel_args.tp_group_) { + const int64_t num_experts = num_total_experts_; + const int64_t intermediate_size = + static_cast(model_args.moe_intermediate_size()); + const std::string& topk_method = model_args.topk_method(); + int64_t ep_size = parallel_args.ep_size(); + int64_t ep_rank = 0; + if (ep_size > 1) { + ep_rank = parallel_args.moe_ep_group_->rank(); + tp_pg_ = parallel_args.moe_tp_group_; + } + + // smoothquant check: If quant_method is not empty, only w8a8 smoothquant is + // supported + if (!quant_args.quant_method().empty()) { + if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 || + !quant_args.activation_dynamic()) { + LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when " + "quant_method is set. " + << "Got quant_method=" << quant_args.quant_method() + << ", bits=" << quant_args.bits() + << ", activation_dynamic=" << quant_args.activation_dynamic(); + } + // If confirmed as smoothquant w8a8, set is_smoothquant_ to true + is_smoothquant_ = true; + } else { + is_smoothquant_ = false; + } + + // calculate the number of experts per rank + num_experts_per_rank_ = num_experts / ep_size; + start_expert_id_ = ep_rank * num_experts_per_rank_; + + if (topk_method == "noaux_tc") { + e_score_correction_bias_ = register_parameter( + "e_score_correction_bias", torch::empty({num_experts}, options), false); + } + + gate_ = register_module( + "gate_proj", + ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options)); + if (n_shared_experts_ > 0) { + /* + The shared_experts are usually implemented using the RowParallelLinear + layer. Typically, this output serves as the enable_result_reduction results + for the module. If only tensor parallelism is applied, immediate + reduction of the shared_experts output isn't necessary; instead, we perform + the reduction once at the end of the MoE operation. + */ + shared_experts_ = + register_module("shared_experts", + DenseMLP(hidden_size_, + intermediate_size * n_shared_experts_, + is_gated_, + false, + hidden_act_, + /*enable_result_reduction=*/false, + quant_args, + tp_pg_, + options)); + shared_expert_gate_ = register_module( + "shared_expert_gate", + torch::nn::Linear( + torch::nn::LinearOptions(hidden_size_, 1).bias(false))); + shared_expert_gate_->weight.set_data( + shared_expert_gate_->weight.to(options)); + } + + // create weight buffer + const int64_t world_size = tp_pg_->world_size(); + int64_t local_intermediate_size = intermediate_size / world_size; + if (is_smoothquant_) { + auto quant_option = options_.dtype(torch::kInt8); + auto fp_option = options_.dtype(torch::kFloat32); + w13_ = register_parameter( + "w13", + torch::empty( + {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_}, + quant_option), + false); + w13_scale_ = register_parameter( + "w13_scale", + torch::empty({num_experts_per_rank_, local_intermediate_size * 2}, + fp_option), + false); + input_smooth_ = register_parameter( + "input_smooth", + torch::empty({num_experts_per_rank_, hidden_size_}, fp_option), + false); + w2_ = register_parameter( + "w2", + torch::empty( + {num_experts_per_rank_, hidden_size_, local_intermediate_size}, + quant_option), + false); + w2_scale_ = register_parameter( + "w2_scale", + torch::empty({num_experts_per_rank_, hidden_size_}, fp_option), + false); + act_smooth_ = register_parameter( + "act_smooth", + torch::empty({num_experts_per_rank_, local_intermediate_size}, + fp_option), + false); + + } else { + w13_ = register_parameter( + "w13", + torch::empty( + {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_}, + options_), + false); + w2_ = register_parameter( + "w2", + torch::empty( + {num_experts_per_rank_, hidden_size_, local_intermediate_size}, + options_), + false); + } +} + +torch::Tensor FusedMoEImpl::select_experts( + const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d, + SelectedExpertInfo& selected_expert_info) { + // prepare the parameters for select_experts + xllm::kernel::MoeFusedTopkParams moe_active_topk_params; + moe_active_topk_params.input = router_logits_2d; + moe_active_topk_params.finished = torch::Tensor(); + moe_active_topk_params.topk = topk_; + moe_active_topk_params.scoring_func = "softmax"; + auto [topk_weights, topk_ids] = + xllm::kernel::moe_active_topk(moe_active_topk_params); + topk_ids = topk_ids.to(torch::kInt32); + if (renormalize_) { + topk_weights = topk_weights / (topk_weights.sum(-1, true) + 1e-6); + } + + xllm::kernel::MoeInitRoutingV2Params moe_init_routing_params; + moe_init_routing_params.x = hidden_states_2d; + moe_init_routing_params.expert_idx = topk_ids; + moe_init_routing_params.scale = std::nullopt; + moe_init_routing_params.offset = std::nullopt; + moe_init_routing_params.active_num = hidden_states_2d.size(0) * topk_; + moe_init_routing_params.expert_capacity = 0; + moe_init_routing_params.expert_num = num_experts_per_rank_; + moe_init_routing_params.drop_pad_mode = 0; + moe_init_routing_params.expert_tokens_num_type = 1; + moe_init_routing_params.expert_tokens_num_flag = true; + moe_init_routing_params.row_idx_type = 0; + std::vector expert_range = { + start_expert_id_, start_expert_id_ + num_experts_per_rank_}; + moe_init_routing_params.active_expert_range = expert_range; + moe_init_routing_params.quant_mode = -1; + // TODO: NPU moe_init_routing_v2 is equivalent to moe_gen_idx + + // moe_expand_input (and the token_count/cusum outputs) on other backends. + auto [expand_hidden_states, expand_row_ids, group_list, dynamic_scale] = + xllm::kernel::moe_init_routing_v2(moe_init_routing_params); + (void)dynamic_scale; + + // collect the selected tensor + selected_expert_info.reduce_weight = topk_weights; + selected_expert_info.combine_idx = expand_row_ids; + selected_expert_info.token_count_slice = group_list; + selected_expert_info.cusum_token_count = group_list; + return expand_hidden_states; +} + +torch::Tensor FusedMoEImpl::forward_expert( + const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + const std::optional& shared_output) { + // prepare the parameters for MoE computation + torch::IntArrayRef hidden_states_shape = hidden_states.sizes(); + torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType(); + torch::Tensor hidden_states_2d = + hidden_states.reshape({-1, hidden_states.size(-1)}); + torch::Tensor router_logits_2d = + router_logits.reshape({-1, router_logits.size(-1)}); + + // Step 1-3: select experts + SelectedExpertInfo selected_expert_info; + torch::Tensor expand_hidden_states = + select_experts(hidden_states_2d, router_logits_2d, selected_expert_info); + + // Step 4: group gemm 1 + torch::Tensor gemm1_out = + create_group_gemm_output(expand_hidden_states, + w13_, + selected_expert_info.token_count_slice, + hidden_states_dtype); + + { + xllm::kernel::GroupGemmParams group_gemm_params; + group_gemm_params.a = expand_hidden_states; + if (w13_.size(1) != expand_hidden_states.size(1)) { + w13_ = w13_.transpose(1, 2); + } + group_gemm_params.b = w13_; + group_gemm_params.group_list = selected_expert_info.token_count_slice; + group_gemm_params.split_item = 2; + group_gemm_params.group_type = 0; + group_gemm_params.group_list_type = 1; + gemm1_out = xllm::kernel::group_gemm(group_gemm_params); + } + + // Step 5: activation + torch::Tensor act_out; + + xllm::kernel::ActivationParams activation_params; + activation_params.input = gemm1_out; + activation_params.output = act_out; + activation_params.act_mode = hidden_act_; + activation_params.is_gated = is_gated_; + xllm::kernel::active(activation_params); + act_out = activation_params.output; + // Step 6: group gemm 2 + torch::Tensor gemm2_out = + create_group_gemm_output(act_out, + w2_, + selected_expert_info.token_count_slice, + hidden_states_dtype); + + { + xllm::kernel::GroupGemmParams group_gemm_params; + group_gemm_params.a = act_out; + if (w2_.size(1) != act_out.size(1)) { + w2_ = w2_.transpose(1, 2); + } + group_gemm_params.b = w2_; + group_gemm_params.group_list = selected_expert_info.token_count_slice; + group_gemm_params.split_item = 2; + group_gemm_params.group_type = 0; + group_gemm_params.group_list_type = 1; + gemm2_out = xllm::kernel::group_gemm(group_gemm_params); + } + + // Step 7: combine the intermediate results and get the final hidden states + torch::Tensor final_hidden_states; + xllm::kernel::MoeCombineResultParams moe_combine_params; + moe_combine_params.input = gemm2_out; + moe_combine_params.reduce_weight = selected_expert_info.reduce_weight; + moe_combine_params.gather_ids = selected_expert_info.combine_idx; + final_hidden_states = xllm::kernel::moe_combine_result(moe_combine_params); + if (shared_output.has_value()) { + final_hidden_states = final_hidden_states + shared_output.value(); + } + // reshape the final hidden states to the original shape + final_hidden_states = final_hidden_states.reshape(hidden_states_shape); + + if (tp_pg_->world_size() > 1) { + final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_); + } + if (parallel_args_.ep_size() > 1) { + final_hidden_states = parallel_state::reduce(final_hidden_states, + parallel_args_.moe_ep_group_); + } + return final_hidden_states; +} + +torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params) { + auto input = hidden_states; + bool need_slice = false; + if (parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1) { + input = parallel_state::gather(input, + parallel_args_.dp_local_process_group_, + input_params.dp_global_token_nums); + need_slice = true; + } + + std::optional shared_output = std::nullopt; + if (n_shared_experts_ > 0) { + shared_output = shared_experts_(input); + if (shared_expert_gate_) { + auto gate = torch::sigmoid(shared_expert_gate_->forward(input)); + if (shared_output.has_value()) { + torch::Tensor res = gate * shared_output.value(); + shared_output = res; + } + } + } + auto router_logits = gate_(input); + auto output = forward_expert(input, router_logits, shared_output); + + if (need_slice) { + const auto& dp_tokens = input_params.dp_global_token_nums; + const int64_t dp_rank = parallel_args_.dp_local_process_group_->rank(); + auto start = + std::accumulate(dp_tokens.begin(), dp_tokens.begin() + dp_rank, 0); + auto end = start + dp_tokens[dp_rank]; + output = output.slice(0, start, end); + } + return output; +} + +void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) { + if (e_score_correction_bias_.defined() && + !e_score_correction_bias_is_loaded_) { + LOAD_WEIGHT(e_score_correction_bias); + } +} + +void FusedMoEImpl::load_experts(const StateDict& state_dict) { + const int64_t rank = tp_pg_->rank(); + const int64_t world_size = tp_pg_->world_size(); + const int64_t start_expert_id = start_expert_id_; + const int64_t num_experts_per_rank = num_experts_per_rank_; + std::vector prefixes = {"gate_proj.", "up_proj."}; + if (is_smoothquant_) { + LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13); + LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale); + LOAD_MOE_WEIGHT("up_proj.", "smooth", input_smooth, -1); + LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1); + LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1); + LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0); + } else { + LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13); + LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1); + + // Some Qwen3.5-MoE checkpoints store expert weights in fused tensors + // (gate_up_proj / down_proj). Fall back to this format when split + // gate_proj/up_proj tensors are absent. + if (!w13_is_loaded_) { + w13_is_loaded_ = load_fused_gate_up_fallback(state_dict, + rank, + world_size, + start_expert_id, + num_experts_per_rank, + w13_); + } + + if (!w2_is_loaded_) { + w2_is_loaded_ = load_fused_down_fallback(state_dict, + rank, + world_size, + start_expert_id, + num_experts_per_rank, + w2_); + } + } +} + +void FusedMoEImpl::load_state_dict(const StateDict& state_dict) { + if (state_dict.size() == 0) { + return; + } + + if (n_shared_experts_ > 0) { + shared_experts_->load_state_dict( + state_dict.get_dict_with_prefix("shared_expert.")); + auto weight = state_dict.get_tensor("shared_expert_gate.weight"); + if (weight.defined()) { + weight = weight.reshape({weight.size(0), -1}); + DCHECK_EQ(shared_expert_gate_->weight.sizes(), weight.sizes()) + << "proj weight size mismatch for " << name(); + shared_expert_gate_->weight.data().copy_(weight); + } + } + + gate_->load_state_dict(state_dict.get_dict_with_prefix("gate.")); + load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate.")); + load_experts(state_dict.get_dict_with_prefix("experts.")); +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.h b/upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.h new file mode 100644 index 00000000..8eb19b60 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.h @@ -0,0 +1,113 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include + +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" +#include "layers/common/dense_mlp.h" +#include "layers/common/fused_moe_base.h" +#include "layers/common/linear.h" + +namespace xllm { +namespace layer { + +class FusedMoEImpl : public torch::nn::Module { + public: + FusedMoEImpl() = default; + FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + torch::Tensor forward_expert( + const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + const std::optional& shared_output); + torch::Tensor forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params); + void load_state_dict(const StateDict& state_dict); + + private: + // struct to store the selected expert info + struct SelectedExpertInfo { + torch::Tensor reduce_weight; + torch::Tensor combine_idx; + torch::Tensor token_count_slice; + torch::Tensor cusum_token_count; + std::optional input_scale; + }; + + // initial steps for MoE computation, select the experts for each token + torch::Tensor select_experts(const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d, + SelectedExpertInfo& selected_expert_info); + + private: + int64_t num_total_experts_; + int64_t topk_; + int64_t num_expert_group_; + int64_t topk_group_; + double route_scale_; + int64_t hidden_size_; + int64_t n_shared_experts_; + bool is_gated_; + bool has_score_bias_; + bool has_bias_; + bool skip_bias_add_; + int64_t renormalize_; + std::string hidden_act_; + std::string scoring_func_; + bool is_smoothquant_; + + int64_t num_experts_per_rank_; + int64_t start_expert_id_; + + ReplicatedLinear gate_{nullptr}; + DenseMLP shared_experts_{nullptr}; + torch::nn::Linear shared_expert_gate_{nullptr}; + QuantArgs quant_args_; + ParallelArgs parallel_args_; + torch::TensorOptions options_; + ProcessGroup* tp_pg_; + + DEFINE_WEIGHT(w13); + DEFINE_FUSED_WEIGHT(w1); + DEFINE_FUSED_WEIGHT(w3); + DEFINE_FUSED_WEIGHT(w2); + DEFINE_WEIGHT(e_score_correction_bias); + DEFINE_WEIGHT(w13_scale); + DEFINE_FUSED_WEIGHT(w1_scale); + DEFINE_FUSED_WEIGHT(w3_scale); + DEFINE_FUSED_WEIGHT(w2_scale); + DEFINE_FUSED_WEIGHT(input_smooth); + DEFINE_FUSED_WEIGHT(act_smooth); + + void load_e_score_correction_bias(const StateDict& state_dict); + void load_experts(const StateDict& state_dict); +}; +TORCH_MODULE(FusedMoE); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.cpp new file mode 100644 index 00000000..c5c9dc7f --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.cpp @@ -0,0 +1,32 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "qwen3_5_decoder_layer_impl.h" + +namespace xllm { +namespace layer { + +Qwen3_5DecoderLayerImpl::Qwen3_5DecoderLayerImpl(const ModelContext& context, + int32_t layer_id) + : Qwen3NextDecoderLayerImpl(context, + layer_id, + std::make_shared( + context.get_model_args(), + context.get_quant_args(), + context.get_parallel_args(), + context.get_tensor_options())) {} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.h new file mode 100644 index 00000000..bb3929d6 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.h @@ -0,0 +1,32 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include "layers/npu_torch/qwen3_5_gated_delta_net.h" +#include "layers/npu_torch/qwen3_next_decoder_layer_impl.h" + +namespace xllm { +namespace layer { + +class Qwen3_5DecoderLayerImpl : public Qwen3NextDecoderLayerImpl { + public: + explicit Qwen3_5DecoderLayerImpl(const ModelContext& context, + int32_t layer_id); +}; +TORCH_MODULE(Qwen3_5DecoderLayer); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.cpp new file mode 100644 index 00000000..9961170a --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.cpp @@ -0,0 +1,240 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "qwen3_next_attention.h" + +#include + +#include +#include + +namespace xllm { +namespace layer { + +Qwen3NextAttentionImpl::Qwen3NextAttentionImpl( + const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options, + int32_t layer_id) { + const int64_t tp_size = parallel_args.tp_group_->world_size(); + const int64_t total_num_heads = args.n_heads(); + const int64_t total_num_kv_heads = args.n_kv_heads().value_or(args.n_heads()); + layer_id_ = layer_id; + rank_ = parallel_args.tp_group_->rank(); + CHECK(total_num_heads % tp_size == 0); + num_heads_ = total_num_heads / tp_size; + + if (total_num_kv_heads >= tp_size) { + CHECK(total_num_kv_heads % tp_size == 0); + num_kv_heads_ = total_num_kv_heads / tp_size; + num_kv_head_replicas_ = 1; + } else { + CHECK(tp_size % total_num_kv_heads == 0); + num_kv_heads_ = 1; + num_kv_head_replicas_ = tp_size / total_num_kv_heads; + } + + head_dim_ = args.head_dim(); + q_size_ = num_heads_ * head_dim_; + kv_size_ = num_kv_heads_ * head_dim_; + scaling_ = 1.0f / std::sqrt(static_cast(head_dim_)); + attn_output_gate_ = args.attn_output_gate(); + // 1. QKV linear + qkv_proj_ = register_module( + "qkv_proj", + QKVParallelLinear(args.hidden_size(), + attn_output_gate_ ? num_heads_ * 2 : num_heads_, + num_kv_heads_, + args.head_dim(), + num_kv_head_replicas_, + /*bias=*/args.attention_bias(), + /*gather_output=*/false, + parallel_args, + options)); + + // 2. O proj + o_proj_ = register_module("o_proj", + RowParallelLinear(total_num_heads * head_dim_, + args.hidden_size(), + /*bias=*/false, + /*input_is_parallelized=*/true, + /*if_reduce_results=*/true, + quant_args, + parallel_args.tp_group_, + options)); + + // 3. Q norm + q_norm_ = register_module( + "q_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options)); + + // 4. K norm + k_norm_ = register_module( + "k_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options)); + + // 5. Rotary embedding + const int rotary_dim = + static_cast(head_dim_ * args.partial_rotary_factor()); + rotary_emb_ = + register_module("rotary_emb", + PartialRotaryEmbedding(rotary_dim, + args.max_position_embeddings(), + args.rope_theta(), + head_dim_, + true, + false, + options)); + + // 6. Attention + attn_ = register_module("attn", + Attention(num_heads_, + head_dim_, + scaling_, + num_kv_heads_, + args.sliding_window())); + + // 7. Fused split_qkv_rmsnorm_mrope kernel setup + rotary_dim_ = static_cast(head_dim_ * args.partial_rotary_factor()); + rms_norm_eps_ = args.rms_norm_eps(); + mrope_section_ = args.rope_scaling_mrope_section(); + is_interleaved_ = args.rope_scaling_mrope_interleaved(); + use_fused_qkv_ = false; + if (attn_output_gate_ && !mrope_section_.empty() && + mrope_section_.size() == 3 && rotary_dim_ > 0 && + xllm::kernel::has_split_qkv_rmsnorm_mrope_specialization( + num_heads_, num_kv_heads_, head_dim_)) { + mrope_gather_pattern_ = + xllm::kernel::build_split_qkv_rmsnorm_mrope_gather_pattern( + rotary_dim_, mrope_section_, is_interleaved_, options.device()); + use_fused_qkv_ = true; + LOG(INFO) << "Qwen3NextAttention layer " << layer_id_ + << ": using fused split_qkv_rmsnorm_mrope kernel"; + } +} + +torch::Tensor Qwen3NextAttentionImpl::build_mrope_cos_sin( + const torch::Tensor& positions) const { + auto cos_sin_cache = rotary_emb_->get_cos_sin_cache(); + if (positions.dim() == 1) { + return cos_sin_cache.index_select(0, positions).repeat({1, 3}); + } + // positions is [3, T] for mRoPE (graph mode or VL) + // transpose from [3, T] to [T, 3] + auto positions_t = positions.permute({1, 0}).contiguous(); + auto gathered = cos_sin_cache.index_select(0, positions_t.view({-1})); + // [T, 3, rope_dim] + return gathered.view({positions.size(1), -1}); +} + +torch::Tensor Qwen3NextAttentionImpl::forward( + const torch::Tensor& positions, + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const torch::Tensor& mrope_cos_sin) { + auto qkv = qkv_proj_->forward(hidden_states); + + if (use_fused_qkv_) { + const int64_t T = qkv.size(0); + xllm::kernel::SplitQkvRmsnormMropeParams params; + params.qkvg = qkv; + params.q_weight = q_norm_->weight(); + params.k_weight = k_norm_->weight(); + params.cos_sin = mrope_cos_sin; + params.gather_pattern = mrope_gather_pattern_; + params.eps = rms_norm_eps_; + params.num_q_heads = num_heads_; + params.num_kv_heads = num_kv_heads_; + params.head_size = head_dim_; + + auto [q, k, v, gate] = xllm::kernel::split_qkv_rmsnorm_mrope(params); + + auto q_flat = q.view({T, q_size_}); + auto k_flat = k.view({T, kv_size_}); + auto v_flat = v.view({T, kv_size_}); + + auto out = std::get<0>( + attn_->forward(attn_metadata, q_flat, k_flat, v_flat, kv_cache)); + out = out * torch::sigmoid(gate.view({T, q_size_})); + return o_proj_->forward(out); + } + + // Fallback path: weight-reordered layout [Q | G | K | V] + torch::Tensor q, k, v; + torch::Tensor gate; + + if (attn_output_gate_) { + q = qkv.slice(-1, 0, q_size_); + gate = qkv.slice(-1, q_size_, q_size_ * 2); + k = qkv.slice(-1, q_size_ * 2, q_size_ * 2 + kv_size_); + v = qkv.slice(-1, q_size_ * 2 + kv_size_, q_size_ * 2 + kv_size_ * 2); + } else { + q = qkv.slice(-1, 0, q_size_); + k = qkv.slice(-1, q_size_, q_size_ + kv_size_); + v = qkv.slice(-1, q_size_ + kv_size_, q_size_ + 2 * kv_size_); + } + + const int64_t T = q.size(0); + auto q_3d = q.view({T, num_heads_, head_dim_}); + q = std::get<0>(q_norm_->forward(q_3d)).view({T, q_size_}); + auto k_3d = k.view({T, num_kv_heads_, head_dim_}); + k = std::get<0>(k_norm_->forward(k_3d)).view({T, kv_size_}); + + rotary_emb_->forward(positions, q, k); + auto out = std::get<0>(attn_->forward(attn_metadata, q, k, v, kv_cache)); + + if (attn_output_gate_) { + out = out * torch::sigmoid(gate); + } + return o_proj_->forward(out); +} + +void Qwen3NextAttentionImpl::load_state_dict(const StateDict& state_dict) { + qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."}); + + if (attn_output_gate_) { + // Rearrange q_proj rows from per-head interleaved [q0,g0,q1,g1,...] + // to grouped [q0,q1,...,g0,g1,...] so forward output is [Q|G|K|V]. + auto w = qkv_proj_->weight(); + auto qg_rows = w.slice(0, 0, q_size_ * 2); + const int64_t hidden = w.size(1); + auto qg_3d = qg_rows.view({num_heads_, 2 * head_dim_, hidden}); + auto q_part = qg_3d.slice(1, 0, head_dim_); + auto g_part = qg_3d.slice(1, head_dim_, 2 * head_dim_); + auto reordered = torch::cat( + {q_part.reshape({q_size_, hidden}), g_part.reshape({q_size_, hidden})}, + 0); + qg_rows.copy_(reordered); + } + + o_proj_->load_state_dict(state_dict.get_dict_with_prefix("o_proj.")); + if (auto w = state_dict.get_tensor("q_norm.weight"); w.defined()) { + q_norm_->load_state_dict(StateDict({{"weight", w}})); + } + if (auto w = state_dict.get_tensor("k_norm.weight"); w.defined()) { + k_norm_->load_state_dict(StateDict({{"weight", w}})); + } + + // Gemma RMSNorm uses (1 + w) as the scale factor, but the fused kernel + // uses standard RMSNorm (w only). Pre-add 1 so the fused kernel produces + // the same result as Qwen3NextRMSNorm (gemma_rms_norm). + if (use_fused_qkv_) { + q_norm_->weight().add_(1.0); + k_norm_->weight().add_(1.0); + } +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.h new file mode 100644 index 00000000..4a4149d1 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.h @@ -0,0 +1,85 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include + +#include "attention.h" +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_args.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "kernels/ops_api.h" +#include "layers/common/linear.h" +#include "layers/common/partial_rotary_embedding.h" +#include "layers/common/qwen3_next_rms_norm.h" + +namespace xllm { +namespace layer { + +class Qwen3NextAttentionImpl : public torch::nn::Module { + public: + Qwen3NextAttentionImpl() = default; + Qwen3NextAttentionImpl(const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options, + int32_t layer_id); + + torch::Tensor forward(const torch::Tensor& positions, + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const torch::Tensor& mrope_cos_sin); + + torch::Tensor build_mrope_cos_sin(const torch::Tensor& positions) const; + + void load_state_dict(const StateDict& state_dict); + + private: + int64_t num_heads_; + int64_t num_kv_heads_; + int64_t num_kv_head_replicas_; + int64_t head_dim_; + int64_t q_size_; + int64_t kv_size_; + float scaling_; + bool attn_output_gate_; + int32_t layer_id_; + int32_t rank_; + int64_t rotary_dim_; + float rms_norm_eps_; + bool use_fused_qkv_; + bool is_interleaved_; + std::vector mrope_section_; + torch::Tensor mrope_gather_pattern_; + + QKVParallelLinear qkv_proj_{nullptr}; + RowParallelLinear o_proj_{nullptr}; + + Qwen3NextRMSNorm q_norm_{nullptr}; + Qwen3NextRMSNorm k_norm_{nullptr}; + + Attention attn_{nullptr}; + PartialRotaryEmbedding rotary_emb_{nullptr}; +}; +TORCH_MODULE(Qwen3NextAttention); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.cpp new file mode 100644 index 00000000..7bb7235b --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.cpp @@ -0,0 +1,41 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "qwen3_next_decoder_layer_impl.h" + +namespace xllm { +namespace layer { + +Qwen3NextDecoderLayerImpl::Qwen3NextDecoderLayerImpl( + const ModelContext& context, + int32_t layer_id) + : Qwen3NextDecoderLayerImpl(context, + layer_id, + std::make_shared( + context.get_model_args(), + context.get_quant_args(), + context.get_parallel_args(), + context.get_tensor_options())) {} + +Qwen3NextDecoderLayerImpl::Qwen3NextDecoderLayerImpl( + const ModelContext& context, + int32_t layer_id, + std::shared_ptr linear_attention_module) + : Qwen3HybridDecoderLayerImplBase(context, + layer_id, + std::move(linear_attention_module)) {} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.h new file mode 100644 index 00000000..b790e1d8 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.h @@ -0,0 +1,38 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include "layers/npu_torch/qwen3_next_gated_delta_net.h" +#include "layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h" + +namespace xllm { +namespace layer { + +class Qwen3NextDecoderLayerImpl : public Qwen3HybridDecoderLayerImplBase { + public: + explicit Qwen3NextDecoderLayerImpl(const ModelContext& context, + int32_t layer_id); + + protected: + Qwen3NextDecoderLayerImpl( + const ModelContext& context, + int32_t layer_id, + std::shared_ptr linear_attention_module); +}; +TORCH_MODULE(Qwen3NextDecoderLayer); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.cpp new file mode 100644 index 00000000..eeeb49f1 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.cpp @@ -0,0 +1,113 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. +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 + https://github.com/jd-opensource/xllm/blob/main/LICENSE +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. +==============================================================================*/ + +#include "qwen3_next_gated_delta_net.h" + +#include + +namespace xllm { +namespace layer { + +Qwen3NextGatedDeltaNetImpl::Qwen3NextGatedDeltaNetImpl( + const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options) + : Qwen3NextGatedDeltaNetImpl(args, + quant_args, + parallel_args, + options, + /*init_projections=*/true) {} + +Qwen3NextGatedDeltaNetImpl::Qwen3NextGatedDeltaNetImpl( + const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options, + bool init_projections) + : Qwen3GatedDeltaNetBaseImpl(args, quant_args, parallel_args, options) { + if (init_projections) { + init_next_projections(args, quant_args, parallel_args, options); + } +} + +void Qwen3NextGatedDeltaNetImpl::init_next_projections( + const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options) { + // QKVZ projection used by Qwen3-Next linear attention. + qkvz_proj_ = register_module("in_proj_qkvz", + ColumnParallelLinear(args.hidden_size(), + k_size_ * 2 + v_size_ * 2, + /*bias=*/false, + /*gather_output=*/false, + quant_args, + parallel_args.tp_group_, + options)); + // BA projection used to derive gating and beta terms. + ba_proj_ = register_module("in_proj_ba", + ColumnParallelLinear(args.hidden_size(), + num_v_heads_ * 2, + /*bias=*/false, + /*gather_output=*/false, + quant_args, + parallel_args.tp_group_, + options)); +} + +std::pair +Qwen3NextGatedDeltaNetImpl::project_padded_inputs( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) { + auto qkvz = qkvz_proj_->forward(hidden_states); + auto ba = ba_proj_->forward(hidden_states); + return {reshape_qkvz_with_pad(attn_metadata, qkvz), + reshape_qkvz_with_pad(attn_metadata, ba)}; +} + +void Qwen3NextGatedDeltaNetImpl::load_state_dict(const StateDict& state_dict) { + load_projection_state_dict(state_dict); + load_common_state_dict(state_dict); +} + +void Qwen3NextGatedDeltaNetImpl::load_projection_state_dict( + const StateDict& state_dict) { + auto qkvz_state_dict = state_dict.get_dict_with_prefix("in_proj_qkvz."); + if (qkvz_state_dict.size() > 0 && !qkvz_proj_->is_weight_loaded()) { + qkvz_proj_->load_state_dict(qkvz_state_dict); + } + + auto ba_state_dict = state_dict.get_dict_with_prefix("in_proj_ba."); + if (ba_state_dict.size() > 0 && !ba_proj_->is_weight_loaded()) { + ba_proj_->load_state_dict(ba_state_dict); + } +} + +void Qwen3NextGatedDeltaNetImpl::verify_loaded_weights( + const std::string& prefix) const { + verify_projection_weights(prefix); + verify_common_loaded_weights(prefix); +} + +void Qwen3NextGatedDeltaNetImpl::verify_projection_weights( + const std::string& prefix) const { + CHECK(qkvz_proj_ && qkvz_proj_->is_weight_loaded()) + << "Missing required weight after all shards loaded: " << prefix + << "in_proj_qkvz.weight"; + CHECK(ba_proj_ && ba_proj_->is_weight_loaded()) + << "Missing required weight after all shards loaded: " << prefix + << "in_proj_ba.weight"; +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.h new file mode 100644 index 00000000..6c8acb9b --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.h @@ -0,0 +1,65 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include +#include + +#include "qwen3_gated_delta_net_base.h" + +namespace xllm { +namespace layer { + +class Qwen3NextGatedDeltaNetImpl : public Qwen3GatedDeltaNetBaseImpl { + public: + Qwen3NextGatedDeltaNetImpl() = default; + Qwen3NextGatedDeltaNetImpl(const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + void load_state_dict(const StateDict& state_dict) override; + void verify_loaded_weights(const std::string& prefix) const override; + + protected: + Qwen3NextGatedDeltaNetImpl(const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options, + bool init_projections); + + std::pair project_padded_inputs( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) override; + + virtual void load_projection_state_dict(const StateDict& state_dict); + virtual void verify_projection_weights(const std::string& prefix) const; + + void init_next_projections(const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + private: + ColumnParallelLinear qkvz_proj_{nullptr}; + ColumnParallelLinear ba_proj_{nullptr}; +}; +TORCH_MODULE(Qwen3NextGatedDeltaNet); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp new file mode 100644 index 00000000..904561d6 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp @@ -0,0 +1,155 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "qwen3_next_hybrid_decoder_layer_base.h" + +#include +#include +#include + +namespace xllm { +namespace layer { + +Qwen3HybridDecoderLayerImplBase::Qwen3HybridDecoderLayerImplBase( + const ModelContext& context, + int32_t layer_id, + std::shared_ptr linear_attention_module) { + const auto& model_args = context.get_model_args(); + const auto& quant_args = context.get_quant_args(); + const auto& parallel_args = context.get_parallel_args(); + const auto& options = context.get_tensor_options(); + const bool use_full_attention = is_full_attention_layer(model_args, layer_id); + + // Initialize attention layers + if (use_full_attention) { + attention_ = register_module( + "self_attn", + Qwen3NextAttention( + model_args, quant_args, parallel_args, options, layer_id)); + } else { + linear_attention_ = + register_module("linear_attn", std::move(linear_attention_module)); + } + + // Initialize norm layers + input_norm_ = register_module( + "input_layernorm", + Qwen3NextRMSNorm( + model_args.hidden_size(), model_args.rms_norm_eps(), options)); + + post_norm_ = register_module( + "post_attention_layernorm", + Qwen3NextRMSNorm( + model_args.hidden_size(), model_args.rms_norm_eps(), options)); + + // Initialize mlp + auto mlp_only_layers = model_args.mlp_only_layers(); + if ((std::count(mlp_only_layers.begin(), mlp_only_layers.end(), layer_id) == + 0) && + model_args.n_routed_experts() > 0 && + (layer_id + 1) % model_args.decoder_sparse_step() == 0) { + moe_mlp_ = register_module("mlp", + FusedMoE(model_args, + FusedMoEArgs{.is_gated = true}, + quant_args, + parallel_args, + options)); + } else { + mlp_ = register_module("mlp", + DenseMLP(model_args.hidden_size(), + model_args.intermediate_size(), + true, + false, + model_args.hidden_act(), + /*enable_result_reduction=*/true, + quant_args, + parallel_args.tp_group_, + options)); + } +} + +void Qwen3HybridDecoderLayerImplBase::load_state_dict( + const StateDict& state_dict) { + if (attention_) { + attention_->load_state_dict(state_dict.get_dict_with_prefix("self_attn.")); + } else { + linear_attention_->load_state_dict( + state_dict.get_dict_with_prefix("linear_attn.")); + } + input_norm_->load_state_dict( + state_dict.get_dict_with_prefix("input_layernorm.")); + post_norm_->load_state_dict( + state_dict.get_dict_with_prefix("post_attention_layernorm.")); + if (moe_mlp_) { + moe_mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp.")); + } else { + mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp.")); + } +} + +void Qwen3HybridDecoderLayerImplBase::verify_loaded_weights( + const std::string& prefix) const { + if (linear_attention_) { + linear_attention_->verify_loaded_weights(prefix + "linear_attn."); + } +} + +torch::Tensor Qwen3HybridDecoderLayerImplBase::forward( + torch::Tensor& x, + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin) { + // Pre-attention norm + if (!residual.has_value()) { + residual = x; + x = std::get<0>(input_norm_->forward(x)); + } else { + std::tie(x, residual) = input_norm_->forward(x, residual); + } + + // Attention + if (attention_) { + x = attention_->forward( + positions, x, attn_metadata, kv_cache, mrope_cos_sin); + } else { + x = linear_attention_->forward(x, attn_metadata, kv_cache, input_params); + } + + // Post-attention norm + std::tie(x, residual) = post_norm_->forward(x, residual); + + // MLP forward + if (moe_mlp_) { + x = moe_mlp_(x, input_params); + } else { + x = mlp_(x); + } + + return x; +} + +torch::Tensor Qwen3HybridDecoderLayerImplBase::build_mrope_cos_sin( + const torch::Tensor& positions) const { + if (attention_) { + return attention_->build_mrope_cos_sin(positions); + } + return {}; +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h new file mode 100644 index 00000000..098f8f2e --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h @@ -0,0 +1,90 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include +#include + +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_input_params.h" +#include "framework/model_context.h" +#include "framework/state_dict/state_dict.h" +#include "layers/common/dense_mlp.h" +#include "layers/common/qwen3_next_rms_norm.h" +#include "layers/npu_torch/fused_moe.h" +#include "layers/npu_torch/qwen3_gated_delta_net_base.h" +#include "layers/npu_torch/qwen3_next_attention.h" + +namespace xllm { +namespace layer { + +class Qwen3HybridDecoderLayerModule : public torch::nn::Module { + public: + virtual void load_state_dict(const StateDict& state_dict) = 0; + virtual void verify_loaded_weights(const std::string& prefix) const = 0; + virtual torch::Tensor forward(torch::Tensor& x, + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin = {}) = 0; + virtual torch::Tensor build_mrope_cos_sin( + const torch::Tensor& positions) const { + return {}; + } +}; + +using Qwen3HybridDecoderLayerModulePtr = + std::shared_ptr; + +class Qwen3HybridDecoderLayerImplBase : public Qwen3HybridDecoderLayerModule { + public: + explicit Qwen3HybridDecoderLayerImplBase( + const ModelContext& context, + int32_t layer_id, + std::shared_ptr linear_attention_module); + + void load_state_dict(const StateDict& state_dict) override; + + void verify_loaded_weights(const std::string& prefix) const override; + + torch::Tensor forward(torch::Tensor& x, + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin = {}) override; + + torch::Tensor build_mrope_cos_sin( + const torch::Tensor& positions) const override; + + protected: + Qwen3NextAttention attention_{nullptr}; + std::shared_ptr linear_attention_; + + DenseMLP mlp_{nullptr}; + FusedMoE moe_mlp_{nullptr}; + + Qwen3NextRMSNorm input_norm_{nullptr}; + Qwen3NextRMSNorm post_norm_{nullptr}; +}; + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_5.h b/upstream_ref/xllm_latest/models/llm/qwen3_5.h new file mode 100644 index 00000000..7c712d39 --- /dev/null +++ b/upstream_ref/xllm_latest/models/llm/qwen3_5.h @@ -0,0 +1,218 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include +#include +#include +#include + +#include "core/layers/npu_torch/qwen3_5_decoder_layer_impl.h" +#include "models/model_registry.h" +#include "qwen3_next.h" + +namespace xllm { + +class Qwen3_5ModelImpl : public Qwen3NextModelImpl { + public: + explicit Qwen3_5ModelImpl(const ModelContext& context) + : Qwen3NextModelImpl(context, /*init_decoder_layers=*/false) { + const int32_t n_layers = context.get_model_args().n_layers(); + for (int32_t layer_id = 0; layer_id < n_layers; ++layer_id) { + add_decoder_layer( + std::make_shared(context, layer_id)); + } + } +}; +TORCH_MODULE(Qwen3_5Model); + +class Qwen3_5ForCausalLMImpl : public Qwen3NextForCausalLMImpl { + public: + explicit Qwen3_5ForCausalLMImpl(const ModelContext& context) + : Qwen3NextForCausalLMImpl(context, /*init_model=*/false) { + set_model_module(std::make_shared(context)); + } +}; +TORCH_MODULE(Qwen3_5ForCausalLM); + +#define LOAD_ARG_TEXT_OR_ROOT(arg_name, json_key, default_value) \ + LOAD_ARG_OR(arg_name, "text_config." json_key, default_value); \ + LOAD_ARG_OR(arg_name, json_key, args->arg_name()) + +#define LOAD_ARG_TEXT_OR_ROOT_CHAIN(arg_name, json_key, default_value) \ + LOAD_ARG_TEXT_OR_ROOT(arg_name, json_key, default_value) + +#define LOAD_QWEN3_5_ROPE_ARG(arg_name, default_value) \ + LOAD_ARG_OR(arg_name, "text_config." #arg_name, default_value); \ + LOAD_ARG_OR(arg_name, #arg_name, args->arg_name()); \ + LOAD_ARG_OR( \ + arg_name, "text_config.rope_scaling." #arg_name, args->arg_name()); \ + LOAD_ARG_OR(arg_name, "rope_scaling." #arg_name, args->arg_name()); \ + LOAD_ARG_OR( \ + arg_name, "text_config.rope_parameters." #arg_name, args->arg_name()); \ + LOAD_ARG_OR(arg_name, "rope_parameters." #arg_name, args->arg_name()) + +#define LOAD_QWEN3_5_NEXT_COMPAT_ARGS(default_moe_intermediate_size, \ + default_num_experts, \ + default_num_experts_per_tok, \ + default_shared_expert_intermediate_size) \ + LOAD_ARG_TEXT_OR_ROOT(attention_bias, "attention_bias", false); \ + LOAD_ARG_TEXT_OR_ROOT(attention_dropout, "attention_dropout", 0.0f); \ + LOAD_ARG_TEXT_OR_ROOT(bos_token_id, "bos_token_id", 151643); \ + LOAD_ARG_TEXT_OR_ROOT(decoder_sparse_step, "decoder_sparse_step", 1); \ + LOAD_ARG_TEXT_OR_ROOT(eos_token_id, "eos_token_id", 151645); \ + LOAD_ARG_TEXT_OR_ROOT(head_dim, "head_dim", 256); \ + LOAD_ARG_TEXT_OR_ROOT(hidden_act, "hidden_act", "silu"); \ + LOAD_ARG_TEXT_OR_ROOT(hidden_size, "hidden_size", 2048); \ + LOAD_ARG_TEXT_OR_ROOT(initializer_range, "initializer_range", 0.02f); \ + LOAD_ARG_TEXT_OR_ROOT(intermediate_size, "intermediate_size", 5120); \ + LOAD_ARG_TEXT_OR_ROOT( \ + max_position_embeddings, "max_position_embeddings", 262144); \ + LOAD_ARG_TEXT_OR_ROOT(max_window_layers, "max_window_layers", 28); \ + LOAD_ARG_TEXT_OR_ROOT(moe_intermediate_size, \ + "moe_intermediate_size", \ + default_moe_intermediate_size); \ + LOAD_ARG_TEXT_OR_ROOT(norm_topk_prob, "norm_topk_prob", true); \ + LOAD_ARG_TEXT_OR_ROOT(n_heads, "num_attention_heads", 16); \ + LOAD_ARG_TEXT_OR_ROOT(num_experts, "num_experts", default_num_experts); \ + LOAD_ARG_TEXT_OR_ROOT(num_experts_per_tok, \ + "num_experts_per_tok", \ + default_num_experts_per_tok); \ + LOAD_ARG_TEXT_OR_ROOT(n_layers, "num_hidden_layers", 48); \ + LOAD_ARG_OR(n_kv_heads, "text_config.num_key_value_heads", 2); \ + LOAD_ARG_OR( \ + n_kv_heads, "num_key_value_heads", args->n_kv_heads().value_or(2)); \ + LOAD_ARG_TEXT_OR_ROOT(output_router_logits, "output_router_logits", false); \ + LOAD_ARG_TEXT_OR_ROOT(rms_norm_eps, "rms_norm_eps", 1e-6); \ + LOAD_QWEN3_5_ROPE_ARG(rope_theta, 10000000.0f); \ + LOAD_ARG_TEXT_OR_ROOT(router_aux_loss_coef, "router_aux_loss_coef", 0.001f); \ + LOAD_ARG_TEXT_OR_ROOT(use_sliding_window, "use_sliding_window", false); \ + LOAD_ARG_TEXT_OR_ROOT(sliding_window, "sliding_window", 4096); \ + LOAD_ARG_TEXT_OR_ROOT(tie_word_embeddings, "tie_word_embeddings", false); \ + LOAD_ARG_TEXT_OR_ROOT(vocab_size, "vocab_size", 151936); \ + LOAD_ARG_TEXT_OR_ROOT( \ + mlp_only_layers, "mlp_only_layers", std::vector()); \ + LOAD_ARG_TEXT_OR_ROOT(attn_output_gate, "attn_output_gate", true); \ + LOAD_ARG_TEXT_OR_ROOT( \ + full_attention_interval, "full_attention_interval", 4); \ + LOAD_ARG_TEXT_OR_ROOT(linear_conv_kernel_dim, "linear_conv_kernel_dim", 4); \ + LOAD_ARG_TEXT_OR_ROOT(linear_key_head_dim, "linear_key_head_dim", 128); \ + LOAD_ARG_TEXT_OR_ROOT(linear_num_key_heads, "linear_num_key_heads", 16); \ + LOAD_ARG_TEXT_OR_ROOT(linear_num_value_heads, "linear_num_value_heads", 32); \ + LOAD_ARG_TEXT_OR_ROOT(linear_value_head_dim, "linear_value_head_dim", 128); \ + LOAD_QWEN3_5_ROPE_ARG(partial_rotary_factor, 0.25f); \ + LOAD_ARG_OR(rope_scaling_mrope_section, \ + "text_config.rope_scaling.mrope_section", \ + std::vector()); \ + LOAD_ARG_OR(rope_scaling_mrope_section, \ + "text_config.rope_parameters.mrope_section", \ + args->rope_scaling_mrope_section()); \ + LOAD_ARG_OR(rope_scaling_mrope_section, \ + "rope_parameters.mrope_section", \ + args->rope_scaling_mrope_section()); \ + LOAD_ARG_OR(rope_scaling_mrope_interleaved, \ + "text_config.rope_scaling.mrope_interleaved", \ + false); \ + LOAD_ARG_OR(rope_scaling_mrope_interleaved, \ + "text_config.rope_parameters.mrope_interleaved", \ + args->rope_scaling_mrope_interleaved()); \ + LOAD_ARG_OR(rope_scaling_mrope_interleaved, \ + "rope_parameters.mrope_interleaved", \ + args->rope_scaling_mrope_interleaved()); \ + LOAD_ARG_TEXT_OR_ROOT(shared_expert_intermediate_size, \ + "shared_expert_intermediate_size", \ + default_shared_expert_intermediate_size); \ + LOAD_ARG_OR( \ + num_nextn_predict_layers, "text_config.mtp_num_hidden_layers", 0); \ + LOAD_ARG_OR(num_nextn_predict_layers, \ + "mtp_num_hidden_layers", \ + args->num_nextn_predict_layers()); \ + LOAD_ARG_OR(num_nextn_predict_layers, \ + "text_config.num_nextn_predict_layers", \ + args->num_nextn_predict_layers()); \ + LOAD_ARG_OR(num_nextn_predict_layers, \ + "num_nextn_predict_layers", \ + args->num_nextn_predict_layers()); \ + LOAD_ARG_OR( \ + layer_types, "text_config.layer_types", std::vector()); \ + LOAD_ARG_OR(layer_types, "layer_types", args->layer_types()); \ + LOAD_ARG_OR( \ + layer_types, "text_config.layers_block_type", args->layer_types()); \ + LOAD_ARG_OR(layer_types, "layers_block_type", args->layer_types()); \ + LOAD_ARG_OR( \ + n_routed_experts, "text_config.n_routed_experts", args->num_experts()); \ + LOAD_ARG_OR(n_routed_experts, "n_routed_experts", args->num_experts()); \ + SET_ARG(n_shared_experts, \ + args->shared_expert_intermediate_size() > 0 ? 1 : 0); \ + SET_ARG(scoring_func, "softmax"); \ + SET_ARG(topk_method, ""); \ + SET_ARG(n_group, -1); \ + SET_ARG(topk_group, 0); \ + SET_ARG(routed_scaling_factor, 1.0f); \ + SET_ARG(stop_token_ids, \ + std::unordered_set({args->eos_token_id()})); \ + LOAD_ARG_TEXT_OR_ROOT(mamba_ssm_dtype, "mamba_ssm_dtype", "float32") + +#define LOAD_QWEN3_5_TYPE_AND_DTYPE(default_model_type) \ + LOAD_ARG_OR(model_type, "model_type", default_model_type); \ + LOAD_ARG_OR(dtype, "text_config.dtype", "bfloat16"); \ + LOAD_ARG_OR(dtype, "dtype", args->dtype()); \ + LOAD_ARG_OR(dtype, "text_config.torch_dtype", args->dtype()); \ + LOAD_ARG_OR(dtype, "torch_dtype", args->dtype()) + +REGISTER_CAUSAL_MODEL(qwen3_5, Qwen3_5ForCausalLM); +REGISTER_MODEL_ARGS(qwen3_5, [&] { + LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5"); + LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/0, + /*num_experts=*/0, + /*num_experts_per_tok=*/0, + /*shared_expert_intermediate_size=*/0); +}); + +REGISTER_CAUSAL_MODEL(qwen3_5_text, Qwen3_5ForCausalLM); +REGISTER_MODEL_ARGS(qwen3_5_text, [&] { + LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_text"); + LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/0, + /*num_experts=*/0, + /*num_experts_per_tok=*/0, + /*shared_expert_intermediate_size=*/0); +}); + +REGISTER_CAUSAL_MODEL(qwen3_5_moe, Qwen3_5ForCausalLM); +REGISTER_MODEL_ARGS(qwen3_5_moe, [&] { + LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_moe"); + LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/512, + /*num_experts=*/512, + /*num_experts_per_tok=*/10, + /*shared_expert_intermediate_size=*/512); +}); + +REGISTER_CAUSAL_MODEL(qwen3_5_moe_text, Qwen3_5ForCausalLM); +REGISTER_MODEL_ARGS(qwen3_5_moe_text, [&] { + LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_moe_text"); + LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/512, + /*num_experts=*/512, + /*num_experts_per_tok=*/10, + /*shared_expert_intermediate_size=*/512); +}); + +#undef LOAD_QWEN3_5_TYPE_AND_DTYPE +#undef LOAD_QWEN3_5_NEXT_COMPAT_ARGS +#undef LOAD_QWEN3_5_ROPE_ARG +#undef LOAD_ARG_TEXT_OR_ROOT_CHAIN +#undef LOAD_ARG_TEXT_OR_ROOT + +} // namespace xllm diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h b/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h new file mode 100644 index 00000000..b4b06b90 --- /dev/null +++ b/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h @@ -0,0 +1,280 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include +#include + +#include "core/framework/model/model_input_params.h" +#include "core/layers/common/linear.h" +#include "models/model_registry.h" +#include "qwen3_5.h" + +namespace xllm { + +namespace { + +StateDict find_lm_head_state_dict(const StateDict& state_dict) { + static const std::vector kLmHeadPrefixes = { + "lm_head.", + "model.lm_head.", + "language_model.lm_head.", + "model.language_model.lm_head."}; + for (const auto& prefix : kLmHeadPrefixes) { + auto sub_dict = state_dict.get_dict_with_prefix(prefix); + if (sub_dict.get_tensor("weight").defined() || + sub_dict.get_tensor("qweight").defined()) { + return sub_dict; + } + } + return StateDict({}, ""); +} + +bool load_qwen3_5_mtp_model_args(const JsonReader& json, + ModelArgs* args, + const std::string& base_model_type, + const std::string& mtp_model_type) { + auto base_loader = ModelRegistry::get_model_args_loader(base_model_type); + if (base_loader == nullptr || base_loader(json, args) == false) { + return false; + } + + int32_t mtp_num_layers = args->num_nextn_predict_layers(); + if (mtp_num_layers <= 0) { + mtp_num_layers = 1; + } + args->model_type(mtp_model_type); + args->num_nextn_predict_layers(mtp_num_layers); + args->n_layers(mtp_num_layers); + args->layer_types(std::vector( + static_cast(mtp_num_layers), "full_attention")); + return true; +} + +} // namespace + +class Qwen3_5MtpModelImpl : public Qwen3HybridModelImplBase { + public: + explicit Qwen3_5MtpModelImpl(const ModelContext& context) + : Qwen3HybridModelImplBase(context) { + const auto& options = context.get_tensor_options(); + const int32_t n_layers = + std::max(static_cast(model_args_.n_layers()), 1); + + pre_fc_norm_embedding_ = register_module( + "pre_fc_norm_embedding", + layer::Qwen3NextRMSNorm( + model_args_.hidden_size(), model_args_.rms_norm_eps(), options)); + pre_fc_norm_hidden_ = register_module( + "pre_fc_norm_hidden", + layer::Qwen3NextRMSNorm( + model_args_.hidden_size(), model_args_.rms_norm_eps(), options)); + fc_ = register_module("fc", + layer::ReplicatedLinear(model_args_.hidden_size() * 2, + model_args_.hidden_size(), + /*bias=*/false, + QuantArgs(), + options)); + + layers_.reserve(n_layers); + for (int32_t layer_id = 0; layer_id < n_layers; ++layer_id) { + add_decoder_layer( + std::make_shared(context, layer_id)); + } + } + + ModelOutput forward(torch::Tensor tokens, + torch::Tensor positions, + std::vector& kv_caches, + const ModelInputParams& input_params) override { + torch::NoGradGuard no_grad; + + if (dp_size_ > 1 && tokens.sizes() == 0) { + tokens = torch::tensor({1}).to(torch::kInt32).to(device_); + positions = torch::tensor({0}).to(torch::kInt32).to(device_); + } + + auto attn_metadata = layer::AttentionMetadataBuilder::build( + input_params, + model_args_.enable_mla(), + build_attention_mask(input_params)); + + torch::Tensor embedding = embed_tokens_(tokens); + torch::Tensor hidden = input_params.input_embedding; + if (hidden.defined() == false) { + hidden = embedding; + } + + embedding = std::get<0>(pre_fc_norm_embedding_->forward(embedding)); + hidden = std::get<0>(pre_fc_norm_hidden_->forward(hidden)); + torch::Tensor mtp_hidden = fc_(torch::cat({embedding, hidden}, -1)); + + CHECK_EQ(kv_caches.size(), layers_.size()); + std::optional residual = std::nullopt; + for (size_t i = 0; i < layers_.size(); ++i) { + mtp_hidden = layers_[i]->forward(mtp_hidden, + residual, + positions, + attn_metadata, + kv_caches[i], + input_params); + } + auto [new_mtp_hidden, new_res] = norm_->forward(mtp_hidden, residual); + mtp_hidden = new_mtp_hidden; + return ModelOutput(mtp_hidden); + } + + void load_state_dict(const StateDict& state_dict) override { + load_shared_embeddings(state_dict); + load_mtp_state_dict(state_dict); + } + + void load_shared_embeddings(const StateDict& state_dict) { + auto embedding_state_dict = + state_dict.get_dict_with_prefix("embed_tokens."); + if (embedding_state_dict.get_tensor("weight").defined()) { + shared_embedding_loaded_ = true; + } + embed_tokens_->load_state_dict(embedding_state_dict); + } + + void load_mtp_state_dict(const StateDict& state_dict) { + if (state_dict.get_tensor("pre_fc_norm_embedding.weight").defined()) { + pre_fc_norm_embedding_loaded_ = true; + } + if (state_dict.get_tensor("pre_fc_norm_hidden.weight").defined()) { + pre_fc_norm_hidden_loaded_ = true; + } + if (state_dict.get_tensor("fc.weight").defined() || + state_dict.get_tensor("fc.qweight").defined()) { + fc_loaded_ = true; + } + if (state_dict.get_tensor("norm.weight").defined()) { + norm_loaded_ = true; + } + + pre_fc_norm_embedding_->load_state_dict( + state_dict.get_dict_with_prefix("pre_fc_norm_embedding.")); + pre_fc_norm_hidden_->load_state_dict( + state_dict.get_dict_with_prefix("pre_fc_norm_hidden.")); + fc_->load_state_dict(state_dict.get_dict_with_prefix("fc.")); + for (size_t i = 0; i < layers_.size(); ++i) { + layers_[i]->load_state_dict( + state_dict.get_dict_with_prefix("layers." + std::to_string(i) + ".")); + } + norm_->load_state_dict(state_dict.get_dict_with_prefix("norm.")); + } + + void verify_loaded_weights(const std::string& prefix) const override { + CHECK(shared_embedding_loaded_) + << "Failed to find shared embedding weights for qwen3.5 mtp draft " + "model"; + CHECK(pre_fc_norm_embedding_loaded_) + << "Failed to find mtp pre_fc_norm_embedding weights for qwen3.5 mtp " + "draft model"; + CHECK(pre_fc_norm_hidden_loaded_) + << "Failed to find mtp pre_fc_norm_hidden weights for qwen3.5 mtp " + "draft model"; + CHECK(fc_loaded_) << "Failed to find mtp fc weights for qwen3.5 mtp draft " + "model"; + CHECK(norm_loaded_) + << "Failed to find mtp norm weights for qwen3.5 mtp draft model"; + for (size_t i = 0; i < layers_.size(); ++i) { + layers_[i]->verify_loaded_weights(prefix + "layers." + std::to_string(i) + + "."); + } + } + + private: + layer::Qwen3NextRMSNorm pre_fc_norm_embedding_{nullptr}; + layer::Qwen3NextRMSNorm pre_fc_norm_hidden_{nullptr}; + layer::ReplicatedLinear fc_{nullptr}; + bool shared_embedding_loaded_ = false; + bool pre_fc_norm_embedding_loaded_ = false; + bool pre_fc_norm_hidden_loaded_ = false; + bool fc_loaded_ = false; + bool norm_loaded_ = false; +}; + +class Qwen3_5MtpForCausalLMImpl : public Qwen3HybridForCausalLMImplBase { + public: + explicit Qwen3_5MtpForCausalLMImpl(const ModelContext& context) + : Qwen3HybridForCausalLMImplBase(context) { + mtp_model_ = std::make_shared(context); + set_model_module(mtp_model_); + } + + void load_model(std::unique_ptr loader) { + static const std::vector kEmbeddingPrefixes = { + "model.language_model.", "language_model.model.", "model.", ""}; + static const std::vector kMtpPrefixes = {"mtp.", "model.mtp."}; + bool lm_head_loaded = false; + + for (const auto& state_dict : loader->get_state_dicts()) { + auto shared_embedding_state_dict = + state_dict->get_dict_with_prefix(kEmbeddingPrefixes); + auto mtp_state_dict = state_dict->get_dict_with_prefix(kMtpPrefixes); + + mtp_model_->load_shared_embeddings(shared_embedding_state_dict); + mtp_model_->load_mtp_state_dict(mtp_state_dict); + + if (tie_word_embeddings_) { + lm_head_->load_state_dict( + shared_embedding_state_dict.get_dict_with_prefix("embed_tokens.")); + if (shared_embedding_state_dict.get_tensor("embed_tokens.weight") + .defined()) { + lm_head_loaded = true; + } + } else { + auto lm_head_state_dict = find_lm_head_state_dict(*state_dict); + lm_head_->load_state_dict(lm_head_state_dict); + if (lm_head_state_dict.get_tensor("weight").defined() || + lm_head_state_dict.get_tensor("qweight").defined()) { + lm_head_loaded = true; + } + } + } + + CHECK(lm_head_loaded) + << "Failed to find lm_head weights for qwen3.5 mtp draft model"; + mtp_model_->verify_loaded_weights("mtp."); + } + + private: + std::shared_ptr mtp_model_; +}; +TORCH_MODULE(Qwen3_5MtpForCausalLM); + +REGISTER_CAUSAL_MODEL(qwen3_5_mtp, Qwen3_5MtpForCausalLM); +REGISTER_CAUSAL_MODEL(qwen3_5_moe_mtp, Qwen3_5MtpForCausalLM); + +REGISTER_MODEL_ARGS_LOADER(qwen3_5_mtp, + [](const JsonReader& json, ModelArgs* args) { + return load_qwen3_5_mtp_model_args( + json, args, "qwen3_5", "qwen3_5_mtp"); + }); + +REGISTER_MODEL_ARGS_LOADER(qwen3_5_moe_mtp, + [](const JsonReader& json, ModelArgs* args) { + return load_qwen3_5_mtp_model_args( + json, args, "qwen3_5_moe", "qwen3_5_moe_mtp"); + }); + +} // namespace xllm diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_next.h b/upstream_ref/xllm_latest/models/llm/qwen3_next.h new file mode 100644 index 00000000..9acf0efc --- /dev/null +++ b/upstream_ref/xllm_latest/models/llm/qwen3_next.h @@ -0,0 +1,126 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include +#include +#include + +#include "core/layers/npu_torch/qwen3_next_decoder_layer_impl.h" +#include "models/model_registry.h" +#include "qwen3_next_hybrid_base.h" + +namespace xllm { + +class Qwen3NextModelImpl : public Qwen3HybridModelImplBase { + public: + explicit Qwen3NextModelImpl(const ModelContext& context) + : Qwen3NextModelImpl(context, /*init_decoder_layers=*/true) {} + + protected: + explicit Qwen3NextModelImpl(const ModelContext& context, + bool init_decoder_layers) + : Qwen3HybridModelImplBase(context) { + if (init_decoder_layers) { + const int32_t n_layers = context.get_model_args().n_layers(); + for (int32_t layer_id = 0; layer_id < n_layers; ++layer_id) { + add_decoder_layer(std::make_shared( + context, layer_id)); + } + } + } +}; +TORCH_MODULE(Qwen3NextModel); + +class Qwen3NextForCausalLMImpl : public Qwen3HybridForCausalLMImplBase { + public: + explicit Qwen3NextForCausalLMImpl(const ModelContext& context) + : Qwen3NextForCausalLMImpl(context, /*init_model=*/true) {} + + protected: + explicit Qwen3NextForCausalLMImpl(const ModelContext& context, + bool init_model) + : Qwen3HybridForCausalLMImplBase(context) { + if (init_model) { + set_model_module(std::make_shared(context)); + } + } +}; +TORCH_MODULE(Qwen3NextForCausalLM); + +// register the causal model +REGISTER_CAUSAL_MODEL(qwen3_next, Qwen3NextForCausalLM); + +// register the model args +REGISTER_MODEL_ARGS(qwen3_next, [&] { + LOAD_ARG_OR(model_type, "model_type", "qwen3_next"); + LOAD_ARG_OR(dtype, "torch_dtype", ""); + LOAD_ARG_OR(attention_bias, "attention_bias", false); + LOAD_ARG_OR(attention_dropout, "attention_dropout", 0.0f); + LOAD_ARG_OR(bos_token_id, "bos_token_id", 151643); + LOAD_ARG_OR(decoder_sparse_step, "decoder_sparse_step", 1); + LOAD_ARG_OR(eos_token_id, "eos_token_id", 151645); + LOAD_ARG_OR(head_dim, "head_dim", 256); + LOAD_ARG_OR(hidden_act, "hidden_act", "silu"); + LOAD_ARG_OR(hidden_size, "hidden_size", 2048); + LOAD_ARG_OR(initializer_range, "initializer_range", 0.02f); + LOAD_ARG_OR(intermediate_size, "intermediate_size", 5120); + LOAD_ARG_OR(max_position_embeddings, "max_position_embeddings", 262144); + LOAD_ARG_OR(max_window_layers, "max_window_layers", 28); + LOAD_ARG_OR(moe_intermediate_size, "moe_intermediate_size", 512); + LOAD_ARG_OR(norm_topk_prob, "norm_topk_prob", true); + LOAD_ARG_OR(n_heads, "num_attention_heads", 16); + LOAD_ARG_OR(num_experts, "num_experts", 512); + LOAD_ARG_OR(num_experts_per_tok, "num_experts_per_tok", 10); + LOAD_ARG_OR(n_layers, "num_hidden_layers", 48); + LOAD_ARG_OR(n_kv_heads, "num_key_value_heads", 2); + LOAD_ARG_OR(output_router_logits, "output_router_logits", false); + LOAD_ARG_OR(rms_norm_eps, "rms_norm_eps", 1e-6); + LOAD_ARG_OR(rope_theta, "rope_theta", 10000000.0f); + LOAD_ARG_OR(router_aux_loss_coef, "router_aux_loss_coef", 0.001f); + LOAD_ARG_OR(use_sliding_window, "use_sliding_window", false); + LOAD_ARG_OR(sliding_window, "sliding_window", 4096); + LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); + LOAD_ARG_OR(vocab_size, "vocab_size", 151936); + LOAD_ARG_OR(mlp_only_layers, "mlp_only_layers", std::vector()); + + // Additional parameters for Qwen3-Next architecture + LOAD_ARG_OR(attn_output_gate, "attn_output_gate", true); + LOAD_ARG_OR(full_attention_interval, "full_attention_interval", 4); + LOAD_ARG_OR(linear_conv_kernel_dim, "linear_conv_kernel_dim", 4); + LOAD_ARG_OR(linear_key_head_dim, "linear_key_head_dim", 128); + LOAD_ARG_OR(linear_num_key_heads, "linear_num_key_heads", 16); + LOAD_ARG_OR(linear_num_value_heads, "linear_num_value_heads", 32); + LOAD_ARG_OR(linear_value_head_dim, "linear_value_head_dim", 128); + LOAD_ARG_OR(partial_rotary_factor, "partial_rotary_factor", 0.25f); + LOAD_ARG_OR( + shared_expert_intermediate_size, "shared_expert_intermediate_size", 512); + LOAD_ARG_OR(layer_types, "layer_types", std::vector()); + + // MoE compatibility with fused_moe implementation. + LOAD_ARG_OR(n_routed_experts, "n_routed_experts", args->num_experts()); + SET_ARG(n_shared_experts, + args->shared_expert_intermediate_size() > 0 ? 1 : 0); + SET_ARG(scoring_func, "softmax"); + SET_ARG(topk_method, ""); + SET_ARG(n_group, -1); + SET_ARG(topk_group, 0); + SET_ARG(routed_scaling_factor, 1.0); + + SET_ARG(stop_token_ids, std::unordered_set({args->eos_token_id()})); +}); + +} // namespace xllm diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h b/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h new file mode 100644 index 00000000..bbf60bc8 --- /dev/null +++ b/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h @@ -0,0 +1,323 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include +#include + +#include "core/common/global_flags.h" +#include "core/framework/kv_cache/kv_cache.h" +#include "core/framework/model/model_input_params.h" +#include "core/framework/model/model_output.h" +#include "core/framework/model_context.h" +#include "core/framework/model_loader.h" +#include "core/layers/common/attention_mask.h" +#include "core/layers/common/attention_metadata_builder.h" +#include "core/layers/common/lm_head.h" +#include "core/layers/common/qwen3_next_rms_norm.h" +#include "core/layers/common/word_embedding.h" +#include "core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h" + +namespace xllm { + +class Qwen3HybridModelModule : public torch::nn::Module { + public: + virtual ModelOutput forward(torch::Tensor tokens, + torch::Tensor positions, + std::vector& kv_caches, + const ModelInputParams& input_params) = 0; + virtual void load_state_dict(const StateDict& state_dict) = 0; + virtual void verify_loaded_weights(const std::string& prefix) const = 0; + virtual layer::WordEmbedding get_word_embedding() = 0; + virtual void set_word_embedding(layer::WordEmbedding& word_embedding) = 0; +}; + +using Qwen3HybridModelModulePtr = std::shared_ptr; + +class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { + public: + explicit Qwen3HybridModelImplBase(const ModelContext& context) + : device_(context.get_tensor_options().device()), + model_args_(context.get_model_args()) { + auto options = context.get_tensor_options(); + auto parallel_args = context.get_parallel_args(); + + blocks_ = register_module("layers", torch::nn::ModuleList()); + layers_.reserve(model_args_.n_layers()); + device_ = options.device(); + dtype_ = options.dtype().toScalarType(); + norm_ = register_module( + "norm", + xllm::layer::Qwen3NextRMSNorm( + model_args_.hidden_size(), model_args_.rms_norm_eps(), options)); + embed_tokens_ = + register_module("embed_tokens", layer::WordEmbedding(context)); + int32_t mask_value = FLAGS_enable_chunked_prefill ? -9984 : 1; + attn_mask_ = layer::AttentionMask(options.device(), + options.dtype().toScalarType(), + /*mask_value=*/mask_value); + dp_size_ = parallel_args.dp_size(); + } + + // tokens: [num_tokens] + // positions: [num_tokens] token pos in the sequence + ModelOutput forward(torch::Tensor tokens, + torch::Tensor positions, + std::vector& kv_caches, + const ModelInputParams& input_params) override { + // Disable gradient computation to reduce memory usage during inference + torch::NoGradGuard no_grad; + if (dp_size_ > 1) { + if (tokens.sizes() == 0) { + tokens = torch::tensor({1}).to(torch::kInt32).to(device_); + positions = torch::tensor({0}).to(torch::kInt32).to(device_); + } + } + + layer::AttentionMetadata attn_metadata = + layer::AttentionMetadataBuilder::build( + input_params, + model_args_.enable_mla(), + build_attention_mask(input_params)); + torch::Tensor h = embed_tokens_(tokens); + + torch::Tensor mrope_cos_sin; + for (const auto& layer : layers_) { + mrope_cos_sin = layer->build_mrope_cos_sin(positions); + if (mrope_cos_sin.defined()) break; + } + + std::optional residual = std::nullopt; + for (size_t i = 0; i < layers_.size(); i++) { + auto& layer = layers_[i]; + h = layer->forward(h, + residual, + positions, + attn_metadata, + kv_caches[i], + input_params, + mrope_cos_sin); + } + auto [hidden_states, residual_out] = norm_->forward(h, residual); + h = hidden_states; + return ModelOutput(h); + } + + // load the weight from the checkpoint + void load_state_dict(const StateDict& state_dict) override { + embed_tokens_->load_state_dict( + state_dict.get_dict_with_prefix("embed_tokens.")); + for (int i = 0; i < static_cast(layers_.size()); i++) { + layers_[i]->load_state_dict( + state_dict.get_dict_with_prefix("layers." + std::to_string(i) + ".")); + } + norm_->load_state_dict(state_dict.get_dict_with_prefix("norm.")); + } + + void verify_loaded_weights(const std::string& prefix) const override { + for (size_t i = 0; i < layers_.size(); ++i) { + layers_[i]->verify_loaded_weights(prefix + "layers." + std::to_string(i) + + "."); + } + } + + layer::WordEmbedding get_word_embedding() override { return embed_tokens_; } + + void set_word_embedding(layer::WordEmbedding& word_embedding) override { + embed_tokens_ = word_embedding; + } + + void add_decoder_layer(layer::Qwen3HybridDecoderLayerModulePtr layer) { + layers_.push_back(layer); + blocks_->push_back(layer); + } + + int32_t num_hidden_layers() const { + return static_cast(layers_.size()); + } + + protected: + torch::Tensor build_attention_mask(const ModelInputParams& input_params) { + max_seq_len_ = std::max(input_params.kv_max_seq_len, max_seq_len_); + if (!FLAGS_enable_chunked_prefill) { + return attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_); + } + + const int32_t num_sequences = input_params.num_sequences; + if (num_sequences <= 0) { + return attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_); + } + + std::vector req_mask_vec; + req_mask_vec.reserve(num_sequences); + for (int32_t j = 0; j < num_sequences; ++j) { + req_mask_vec.emplace_back( + attn_mask_.gen_append_mask(input_params.q_seq_lens_vec[j], + input_params.kv_seq_lens_vec[j], + max_seq_len_, + dtype_, + device_)); + } + return torch::cat(req_mask_vec, 0); + } + + ModelArgs model_args_; + torch::nn::ModuleList blocks_{nullptr}; + std::vector layers_; + int32_t max_seq_len_ = 0; + int32_t dp_size_ = 1; + torch::Device device_; + torch::ScalarType dtype_ = torch::kFloat; + layer::Qwen3NextRMSNorm norm_{nullptr}; + layer::AttentionMask attn_mask_; + layer::WordEmbedding embed_tokens_{nullptr}; +}; + +class Qwen3HybridForCausalLMImplBase : public torch::nn::Module { + public: + explicit Qwen3HybridForCausalLMImplBase(const ModelContext& context) { + tie_word_embeddings_ = context.get_model_args().tie_word_embeddings(); + lm_head_ = register_module("lm_head", layer::LmHead(context)); + } + + // tokens: [num_tokens] + // positions: [num_tokens] token pos in the sequence + // returns: [num_tokens, hidden_size] + ModelOutput forward(const torch::Tensor& tokens, + const torch::Tensor& positions, + std::vector& kv_caches, + const ModelInputParams& input_params) { + return model_->forward(tokens, positions, kv_caches, input_params); + } + + // hidden_states: [num_tokens, hidden_size] + // seleted_idxes: [num_tokens] + // returns: [num_tokens, vocab_size] + torch::Tensor logits(const torch::Tensor& hidden_states, + const torch::Tensor& seleted_idxes) { + auto h = hidden_states; + if (seleted_idxes.defined()) { + h = h.index_select(/*dim=*/0, seleted_idxes); + } + return lm_head_(h); + } + + // hidden_states: [num_tokens, hidden_size] + // seleted_idxes: [num_tokens] + torch::Tensor pooler(const torch::Tensor& hidden_states, + const torch::Tensor& seleted_idxes) { + auto h = hidden_states; + if (seleted_idxes.defined()) { + h = h.index_select(/*dim=*/0, seleted_idxes); + } + namespace F = torch::nn::functional; + return F::normalize(h, F::NormalizeFuncOptions().p(2).dim(1)); + } + + void load_model(std::unique_ptr loader) { + auto has_model_weights = [](const StateDict& dict) { + return dict.get_tensor("embed_tokens.weight").defined() || + dict.get_dict_with_prefix("layers.").size() > 0 || + dict.get_tensor("norm.weight").defined(); + }; + auto has_lm_head_weights = [](const StateDict& dict) { + return dict.get_tensor("weight").defined() || + dict.get_tensor("qweight").defined(); + }; + + for (const auto& state_dict : loader->get_state_dicts()) { + auto model_state_dict = state_dict->get_dict_with_prefix("model."); + if (!has_model_weights(model_state_dict)) { + auto language_model_state_dict = + state_dict->get_dict_with_prefix("language_model.model."); + if (has_model_weights(language_model_state_dict)) { + model_state_dict = language_model_state_dict; + } else { + auto wrapped_language_model_state_dict = + state_dict->get_dict_with_prefix("model.language_model."); + if (has_model_weights(wrapped_language_model_state_dict)) { + model_state_dict = wrapped_language_model_state_dict; + } + } + } + model_->load_state_dict(model_state_dict); + + auto lm_head_state_dict = state_dict->get_dict_with_prefix("lm_head."); + if (!has_lm_head_weights(lm_head_state_dict)) { + auto language_model_lm_head_state_dict = + state_dict->get_dict_with_prefix("language_model.lm_head."); + if (has_lm_head_weights(language_model_lm_head_state_dict)) { + lm_head_state_dict = language_model_lm_head_state_dict; + } else { + auto wrapped_language_model_lm_head_state_dict = + state_dict->get_dict_with_prefix("model.language_model.lm_head."); + if (has_lm_head_weights(wrapped_language_model_lm_head_state_dict)) { + lm_head_state_dict = wrapped_language_model_lm_head_state_dict; + } else { + auto wrapped_lm_head_state_dict = + state_dict->get_dict_with_prefix("model.lm_head."); + if (has_lm_head_weights(wrapped_lm_head_state_dict)) { + lm_head_state_dict = wrapped_lm_head_state_dict; + } + } + } + } + if (!has_lm_head_weights(lm_head_state_dict) && tie_word_embeddings_) { + auto tied_lm_head_state_dict = + model_state_dict.get_dict_with_prefix("embed_tokens."); + if (has_lm_head_weights(tied_lm_head_state_dict)) { + lm_head_state_dict = tied_lm_head_state_dict; + } + } + lm_head_->load_state_dict(lm_head_state_dict); + } + + model_->verify_loaded_weights("model."); + } + + virtual void prepare_expert_weight(int32_t layer_id, + const std::vector& expert_ids) { + return; + } + virtual void update_expert_weight(int32_t layer_id) { return; } + + layer::LmHead get_lm_head() { return lm_head_; } + + void set_lm_head(layer::LmHead& head) { lm_head_ = head; } + + layer::WordEmbedding get_word_embedding() { + return model_->get_word_embedding(); + } + + void set_word_embedding(layer::WordEmbedding& word_embedding) { + model_->set_word_embedding(word_embedding); + } + + void set_model_module(Qwen3HybridModelModulePtr model) { + model_ = register_module("model", std::move(model)); + } + + protected: + bool tie_word_embeddings_{false}; + layer::LmHead lm_head_{nullptr}; + Qwen3HybridModelModulePtr model_; +}; + +} // namespace xllm diff --git a/upstream_ref/xllm_latest/models/vlm/qwen3_5.h b/upstream_ref/xllm_latest/models/vlm/qwen3_5.h new file mode 100644 index 00000000..9a9a9c02 --- /dev/null +++ b/upstream_ref/xllm_latest/models/vlm/qwen3_5.h @@ -0,0 +1,312 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include "core/framework/model/model_output.h" +#include "core/layers/common/lm_head.h" +#include "core/layers/common/qwen3_next_rms_norm.h" +#include "core/layers/common/rms_norm.h" +#include "core/layers/mlu/qwen3_5_decoder_layer.h" +#include "core/layers/qwen3_vision_layer.h" +#include "models/llm/llm_model_base.h" +#include "models/model_registry.h" +#include "models/vlm/qwen3_vl_base.h" +#include "processors/input_processor.h" +#include "processors/qwen2_vl_image_processor.h" +#include "qwen3_vl.h" + +namespace xllm { +class Qwen3_5ModelImpl final + : public LlmModelImplBase { + public: + Qwen3_5ModelImpl(const ModelContext& context) + : LlmModelImplBase("qwen3_5", + context.get_model_args()) { + auto model_args = context.get_model_args(); + auto options = context.get_tensor_options(); + auto parallel_args = context.get_parallel_args(); + dp_size_ = parallel_args.dp_size(); + + if (!mrope_section_.empty()) { + int64_t rotary_dim = static_cast( + model_args.head_dim() * model_args.partial_rotary_factor()); + cos_sin_ = layer::rotary::get_concat_rotary_embedding( + rotary_dim, + model_args.max_position_embeddings(), + model_args.rope_theta(), + options); + } + + layers_.reserve(model_args.n_layers()); + rms_norm_ = register_module( + "norm", + layer::Qwen3NextRMSNorm( + model_args.hidden_size(), model_args.rms_norm_eps(), options)); + embed_tokens_ = + register_module("embed_tokens", layer::WordEmbedding(context)); + + for (int32_t i = 0; i < model_args.n_layers(); i++) { + auto layer = layer::Qwen3_5DecoderLayer(context, i); + layers_.push_back(layer); + } + } + + void load_state_dict(const StateDict& state_dict) override { + embed_tokens_->load_state_dict( + state_dict.get_dict_with_prefix("embed_tokens.")); + + // call each layer's load_state_dict function + for (size_t i = 0; i < layers_.size(); i++) { + layers_[i]->load_state_dict( + state_dict.get_dict_with_prefix("layers." + std::to_string(i) + ".")); + } + rms_norm_->load_state_dict(state_dict.get_dict_with_prefix("norm.")); + } + + std::pair apply_mrope( + const torch::Tensor positions) override { + auto target_cos_sin = cos_sin_.index({positions}); + auto target_cos_sin_chunks = target_cos_sin.chunk(/*chunks=*/2, /*dim=*/-1); + auto cos_pos = target_cos_sin_chunks[0].contiguous(); + auto sin_pos = target_cos_sin_chunks[1].contiguous(); + auto apply = [this](torch::Tensor x) { + auto freqs_t = x[0].clone(); + int64_t mrop_length = static_cast(freqs_t.size(-1) / 2); + + for (int32_t dim_idx = 1; dim_idx <= 2; ++dim_idx) { + int64_t offset = dim_idx; + int64_t section_len = mrope_section_[dim_idx]; + int64_t length = section_len * 3; + + auto idx_first_half = torch::arange(offset, length, 3, torch::kLong); + auto idx_second_half = torch::arange( + offset + mrop_length, length + mrop_length, 3, torch::kLong); + + auto idx_tensor = + torch::cat({idx_first_half, idx_second_half}, 0).to(x.device()); + auto src = x[dim_idx].index_select(-1, idx_tensor); + freqs_t.index_copy_(-1, idx_tensor, src); + } + return freqs_t; + }; + cos_pos = apply(cos_pos.reshape({positions.size(0), -1, cos_pos.size(-1)})); + sin_pos = apply(sin_pos.reshape({positions.size(0), -1, sin_pos.size(-1)})); + return std::make_pair(cos_pos, sin_pos); + } + + virtual ModelOutput forward(torch::Tensor tokens, + torch::Tensor positions, + std::vector& kv_caches, + const ModelInputParams& input_params) { + ModelInputParams& input_params_new = + const_cast(input_params); + std::vector deep_stacks; + + if (dp_size_ > 1) { + if (tokens.numel() == 0) { + tokens = torch::tensor({1}).to(torch::kInt32).to(tokens.device()); + positions = torch::tensor({1}).to(torch::kInt32).to(positions.device()); + } + auto& dp_token_nums = input_params_new.dp_global_token_nums; + std::replace(dp_token_nums.begin(), dp_token_nums.end(), 0, 1); + } + + auto inputs_embeds = input_params.input_embedding; + torch::Tensor h; + if (inputs_embeds.defined()) { + h = inputs_embeds; + } else { + h = embed_tokens_(tokens); + } + + if (!input_params_new.attn_metadata) { + input_params_new.attn_metadata = + std::make_shared( + get_attention_metadata(input_params_new, h)); + } + + auto& attn_metadata = *(input_params_new.attn_metadata); + bool only_prefill = + (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill); + if (positions.dim() == 2 && only_prefill && !mrope_section_.empty()) { + std::tie(attn_metadata.mrope_cos, attn_metadata.mrope_sin) = + apply_mrope(positions); + } + + std::optional residual; + for (size_t i = 0; i < layers_.size(); i++) { + auto& layer = layers_[i]; + h = layer(h, + residual, + positions, + attn_metadata, + kv_caches[i], + input_params_new); + } + if (residual.has_value()) { + h = h + residual.value(); + } + auto hidden_states = std::get<0>(rms_norm_(h)); + return ModelOutput(hidden_states); + } + + private: + int32_t dp_size_ = 1; + layer::Qwen3NextRMSNorm rms_norm_{nullptr}; + layer::AttentionMetadata get_attention_metadata( + const ModelInputParams& params, + const torch::Tensor& h) { + auto attn_metadata = layer::AttentionMetadataBuilder::build(params, false); + // TODO: support linear attention + return attn_metadata; + } +}; +TORCH_MODULE(Qwen3_5Model); + +class Qwen3_5ForCausalLMImpl : public LlmForCausalLMImplBase { + public: + Qwen3_5ForCausalLMImpl(const ModelContext& context) + : LlmForCausalLMImplBase(context) {} + + torch::Tensor pooler(const torch::Tensor& hidden_states, + const torch::Tensor& seleted_idxes) { + auto h = hidden_states; + if (seleted_idxes.defined()) { + h = h.index_select(/*dim=*/0, seleted_idxes); + } + namespace F = torch::nn::functional; + return F::normalize(h, F::NormalizeFuncOptions().p(2).dim(1)); + } +}; +TORCH_MODULE(Qwen3_5ForCausalLM); + +using Qwen3_5ForConditionalGenerationImpl = + Qwen3VLForConditionalGenerationBase; +TORCH_MODULE(Qwen3_5ForConditionalGeneration); + +#define LOAD_QWEN3_5_COMMON_ARGS() \ + LOAD_ARG_OR(model_type, "model_type", "qwen3_5"); \ + LOAD_ARG_OR(dtype, "text_config.dtype", "bfloat16"); \ + LOAD_ARG_OR(vocab_size, "text_config.vocab_size", 248320); \ + LOAD_ARG_OR(hidden_size, "text_config.hidden_size", 5120); \ + LOAD_ARG_OR(hidden_act, "text_config.hidden_act", "silu"); \ + LOAD_ARG_OR(intermediate_size, "text_config.intermediate_size", 17408); \ + LOAD_ARG_OR(n_layers, "text_config.num_hidden_layers", 64); \ + LOAD_ARG_OR(n_heads, "text_config.num_attention_heads", 24); \ + LOAD_ARG(n_kv_heads, "text_config.num_key_value_heads"); \ + LOAD_ARG_OR( \ + max_position_embeddings, "text_config.max_position_embeddings", 262144); \ + LOAD_ARG_OR(rms_norm_eps, "text_config.rms_norm_eps", 1e-6); \ + LOAD_ARG_OR(eos_token_id, "text_config.eos_token_id", 248044); \ + LOAD_ARG_OR( \ + rope_theta, "text_config.rope_parameters.rope_theta", 10000000.0f); \ + LOAD_ARG_OR(head_dim, "text_config.head_dim", 256); \ + LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); \ + LOAD_ARG(layer_types, "text_config.layer_types"); \ + LOAD_ARG_OR( \ + linear_conv_kernel_dim, "text_config.linear_conv_kernel_dim", 4); \ + LOAD_ARG_OR(linear_key_head_dim, "text_config.linear_key_head_dim", 128); \ + LOAD_ARG_OR( \ + linear_value_head_dim, "text_config.linear_value_head_dim", 128); \ + LOAD_ARG_OR(linear_num_key_heads, "text_config.linear_num_key_heads", 16); \ + LOAD_ARG_OR( \ + linear_num_value_heads, "text_config.linear_num_value_heads", 48); \ + LOAD_ARG_OR( \ + full_attention_interval, "text_config.full_attention_interval", 4); \ + LOAD_ARG_OR(attn_output_gate, "text_config.attn_output_gate", false); \ + LOAD_ARG_OR( \ + num_nextn_predict_layers, "text_config.mtp_num_hidden_layers", 0); \ + LOAD_ARG_OR(attention_bias, "text_config.attention_bias", false); \ + LOAD_ARG_OR(attention_dropout, "text_config.attention_dropout", 0.0f); \ + LOAD_ARG_OR(initializer_range, "text_config.initializer_range", 0.02f); \ + LOAD_ARG_OR( \ + mlp_only_layers, "text_config.mlp_only_layers", std::vector()); \ + LOAD_ARG(rope_scaling_mrope_section, \ + "text_config.rope_parameters.mrope_section"); \ + LOAD_ARG_OR(rope_scaling_rope_type, \ + "text_config.rope_parameters.rope_type", \ + "default"); \ + LOAD_ARG_OR(partial_rotary_factor, \ + "text_config.rope_parameters.partial_rotary_factor", \ + 0.25f) + +#define LOAD_QWEN3_5_VISION_ARGS() \ + LOAD_ARG_OR(image_token_id, "image_token_id", 248056); \ + LOAD_ARG_OR(video_token_id, "video_token_id", 248057); \ + LOAD_ARG_OR(vision_start_token_id, "vision_start_token_id", 248053); \ + LOAD_ARG_OR(vision_end_token_id, "vision_end_token_id", 248054); \ + LOAD_ARG(mm_deepstack_visual_indexes, \ + "vision_config.deepstack_visual_indexes"); \ + LOAD_ARG_OR(mm_num_hidden_layers, "vision_config.depth", 27); \ + LOAD_ARG_OR(mm_hidden_act, "vision_config.hidden_act", "gelu_pytorch_tanh"); \ + LOAD_ARG_OR(mm_hidden_size, "vision_config.hidden_size", 1152); \ + LOAD_ARG_OR(mm_num_channels, "vision_config.in_channels", 3); \ + LOAD_ARG_OR(mm_initializer_range, "vision_config.initializer_range", 0.02f); \ + LOAD_ARG_OR(mm_intermediate_size, "vision_config.intermediate_size", 4304); \ + LOAD_ARG_OR(mm_num_attention_heads, "vision_config.num_heads", 16); \ + LOAD_ARG_OR(mm_num_position_embeddings, \ + "vision_config.num_position_embeddings", \ + 2304); \ + LOAD_ARG_OR(mm_projection_dim, "vision_config.out_hidden_size", 5120); \ + LOAD_ARG_OR(mm_patch_size, "vision_config.patch_size", 16); \ + LOAD_ARG_OR(mm_spatial_merge_size, "vision_config.spatial_merge_size", 2); \ + LOAD_ARG_OR(mm_temporal_patch_size, "vision_config.temporal_patch_size", 2); \ + LOAD_ARG_OR_FUNC(mm_head_dim, "head_dim", [&] { \ + return args->mm_hidden_size() / args->mm_num_attention_heads(); \ + }); \ + LOAD_ARG_OR( \ + rope_scaling_rope_type, "vision_config.rope_scaling.type", "mrope") + +REGISTER_INPUT_PROCESSOR(qwen3_5, Qwen2_5_VLInputProcessor); +REGISTER_CAUSAL_VLM_MODEL(qwen3_5, Qwen3_5ForConditionalGeneration); +REGISTER_IMAGE_PROCESSOR(qwen3_5, Qwen2VLImageProcessor); +REGISTER_MODEL_ARGS(qwen3_5, [&] { + LOAD_QWEN3_5_COMMON_ARGS(); + LOAD_QWEN3_5_VISION_ARGS(); + SET_ARG(stop_token_ids, std::unordered_set({args->eos_token_id()})); +}); + +REGISTER_INPUT_PROCESSOR(qwen3_5_moe, Qwen2_5_VLInputProcessor); +REGISTER_CAUSAL_VLM_MODEL(qwen3_5_moe, Qwen3_5ForConditionalGeneration); +REGISTER_IMAGE_PROCESSOR(qwen3_5_moe, Qwen2VLImageProcessor); +REGISTER_MODEL_ARGS(qwen3_5_moe, [&] { + LOAD_QWEN3_5_COMMON_ARGS(); + LOAD_QWEN3_5_VISION_ARGS(); + + LOAD_ARG_OR(decoder_sparse_step, "text_config.decoder_sparse_step", 1); + LOAD_ARG_OR(moe_intermediate_size, "text_config.moe_intermediate_size", 512); + LOAD_ARG_OR(num_experts, "text_config.num_experts", 512); + LOAD_ARG_OR(num_experts_per_tok, "text_config.num_experts_per_tok", 10); + LOAD_ARG_OR(shared_expert_intermediate_size, + "text_config.shared_expert_intermediate_size", + 512); + LOAD_ARG_OR(norm_topk_prob, "text_config.norm_topk_prob", true); + + LOAD_ARG_OR( + n_routed_experts, "text_config.n_routed_experts", args->num_experts()); + SET_ARG(n_shared_experts, + args->shared_expert_intermediate_size() > 0 ? 1 : 0); + SET_ARG(scoring_func, "softmax"); + SET_ARG(topk_method, ""); + SET_ARG(n_group, -1); + SET_ARG(topk_group, 0); + SET_ARG(routed_scaling_factor, 1.0f); + + SET_ARG(stop_token_ids, std::unordered_set({args->eos_token_id()})); +}); + +} // namespace xllm