Upload 14 files

Long_Term_Forecasting/data_loader.py

@@ -0,0 +1,314 @@
+import os
+import numpy as np
+import pandas as pd
+from torch.utils.data import Dataset
+from sklearn.preprocessing import StandardScaler
+
+
+
+import warnings
+
+warnings.filterwarnings('ignore')
+
+
+class Dataset_ETT_hour(Dataset):
+    def __init__(self, root_path, flag='train', size=None,
+                 features='S', data_path='ETTh1.csv',
+                 target='OT', scale=True, timeenc=0, freq='h', seasonal_patterns=None):
+        # size [seq_len, label_len, pred_len]
+        # info
+        if size == None:
+            self.seq_len = 24 * 4 * 4
+            self.label_len = 24 * 4
+            self.pred_len = 24 * 4
+        else:
+            self.seq_len = size[0]
+            self.label_len = size[1]
+            self.pred_len = size[2]
+        # init
+        assert flag in ['train', 'test', 'val']
+        type_map = {'train': 0, 'val': 1, 'test': 2}
+        self.set_type = type_map[flag]
+
+        self.features = features
+        self.target = target
+        self.scale = scale
+        self.timeenc = timeenc
+        self.freq = freq
+
+        self.root_path = root_path
+        self.data_path = data_path
+        self.raw_start = 96
+        self.raw_last= 720
+        self.seperate = [96, 192, 336, 720]
+        self.__read_data__()
+
+    def __read_data__(self):
+        self.scaler = StandardScaler()
+        df_raw = pd.read_csv(os.path.join(self.root_path,
+                                          self.data_path))
+
+        border1s = [0, 12 * 30 * 24 - self.seq_len, 12 * 30 * 24 + 4 * 30 * 24 - self.seq_len]
+        border2s = [12 * 30 * 24, 12 * 30 * 24 + 4 * 30 * 24, 12 * 30 * 24 + 8 * 30 * 24]
+        border1 = border1s[self.set_type]
+        border2 = border2s[self.set_type]
+
+        if self.features == 'M' or self.features == 'MS':
+            cols_data = df_raw.columns[1:]
+            df_data = df_raw[cols_data]
+        elif self.features == 'S':
+            df_data = df_raw[[self.target]]
+
+        if self.scale:
+            train_data = df_data[border1s[0]:border2s[0]]
+            self.scaler.fit(train_data.values)
+            data = self.scaler.transform(df_data.values)
+        else:
+            data = df_data.values
+
+        df_stamp = df_raw[['date']][border1:border2]
+        df_stamp['date'] = pd.to_datetime(df_stamp.date)
+        if self.timeenc == 0:
+            df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
+            df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
+            df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
+            df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
+
+
+
+        self.data_x = data[border1:border2]
+        self.data_y = data[border1:border2]
+
+
+    def __getitem__(self, index):
+        s_begin = index
+        s_end = s_begin + self.seq_len
+        r_begin = s_end - self.label_len
+        r_end = r_begin + self.label_len + self.pred_len
+              
+        seq_x = self.data_x[s_begin:s_end]
+        seq_y = self.data_y[r_begin:r_end]
+        res = []
+        for end in self.seperate:
+            r_end = r_begin + self.label_len + end
+            
+            if r_end <= self.data_y.shape[0]:
+                res.append(self.data_y[r_begin:r_end])
+            else:
+                res.append(np.full((r_end - r_begin,self.data_y.shape[-1]), np.nan))
+
+        return seq_x, seq_y, res[0], res[1], res[2], res[3]
+
+    def __len__(self):
+        return len(self.data_x) - self.seq_len - self.raw_start + 1
+
+    def inverse_transform(self, data):
+        return self.scaler.inverse_transform(data)
+
+
+class Dataset_ETT_minute(Dataset):
+    def __init__(self, root_path, flag='train', size=None,
+                 features='S', data_path='ETTm1.csv',
+                 target='OT', scale=True, timeenc=0, freq='t', seasonal_patterns=None):
+        # size [seq_len, label_len, pred_len]
+        # info
+        if size == None:
+            self.seq_len = 24 * 4 * 4
+            self.label_len = 24 * 4
+            self.pred_len = 24 * 4
+        else:
+            self.seq_len = size[0]
+            self.label_len = size[1]
+            self.pred_len = size[2]
+        # init
+        assert flag in ['train', 'test', 'val']
+        type_map = {'train': 0, 'val': 1, 'test': 2}
+        self.set_type = type_map[flag]
+
+        self.features = features
+        self.target = target
+        self.scale = scale
+        self.timeenc = timeenc
+        self.freq = freq
+
+        self.root_path = root_path
+        self.data_path = data_path
+        
+        self.raw_start = 96
+        self.raw_last= 720
+        self.seperate = [96, 192, 336, 720]
+        self.__read_data__()
+
+    def __read_data__(self):
+        self.scaler = StandardScaler()
+        df_raw = pd.read_csv(os.path.join(self.root_path,
+                                          self.data_path))
+
+        border1s = [0, 12 * 30 * 24 * 4 - self.seq_len, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4 - self.seq_len]
+        border2s = [12 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 8 * 30 * 24 * 4]
+        border1 = border1s[self.set_type]
+        border2 = border2s[self.set_type]
+
+        if self.features == 'M' or self.features == 'MS':
+            cols_data = df_raw.columns[1:]
+            df_data = df_raw[cols_data]
+        elif self.features == 'S':
+            df_data = df_raw[[self.target]]
+
+        if self.scale:
+            train_data = df_data[border1s[0]:border2s[0]]
+            self.scaler.fit(train_data.values)
+            data = self.scaler.transform(df_data.values)
+        else:
+            data = df_data.values
+
+        df_stamp = df_raw[['date']][border1:border2]
+        df_stamp['date'] = pd.to_datetime(df_stamp.date)
+        if self.timeenc == 0:
+            df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
+            df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
+            df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
+            df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
+            df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)
+            df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)
+
+
+
+        self.data_x = data[border1:border2]
+        self.data_y = data[border1:border2]
+
+    def __getitem__(self, index):
+
+        s_begin = index
+        s_end = s_begin + self.seq_len
+        r_begin = s_end - self.label_len
+        r_end = r_begin + self.label_len + self.pred_len
+
+        seq_x = self.data_x[s_begin:s_end]
+        seq_y = self.data_y[r_begin:r_end]
+        res = []
+        for end in self.seperate:
+            r_end = r_begin + self.label_len + end
+
+            if r_end <= self.data_y.shape[0]:
+                res.append(self.data_y[r_begin:r_end])
+            else:
+                res.append(np.full((r_end - r_begin,self.data_y.shape[-1]), np.nan))
+
+        return seq_x, seq_y, res[0], res[1], res[2], res[3]
+
+    def __len__(self):
+        return len(self.data_x) - self.seq_len - self.raw_start + 1
+
+    def inverse_transform(self, data):
+        return self.scaler.inverse_transform(data)
+
+
+class Dataset_Custom(Dataset):
+    def __init__(self, root_path, flag='train', size=None,
+                 features='S', data_path='ETTh1.csv',
+                 target='OT', scale=True, timeenc=0, freq='h', seasonal_patterns=None):
+        # size [seq_len, label_len, pred_len]
+        # info
+        if size == None:
+            self.seq_len = 24 * 4 * 4
+            self.label_len = 24 * 4
+            self.pred_len = 24 * 4
+        else:
+            self.seq_len = size[0]
+            self.label_len = size[1]
+            self.pred_len = size[2]
+        # init
+        assert flag in ['train', 'test', 'val']
+        type_map = {'train': 0, 'val': 1, 'test': 2}
+        self.set_type = type_map[flag]
+
+        self.features = features
+        self.target = target
+        self.scale = scale
+        self.timeenc = timeenc
+        self.freq = freq
+
+        self.root_path = root_path
+        self.data_path = data_path
+        self.raw_start = 96
+        self.raw_last= 720
+        self.seperate = [96, 192, 336, 720]
+        self.__read_data__()
+
+    def __read_data__(self):
+        self.scaler = StandardScaler()
+        df_raw = pd.read_csv(os.path.join(self.root_path,
+                                          self.data_path))
+
+        '''
+        df_raw.columns: ['date', ...(other features), target feature]
+        '''
+        cols = list(df_raw.columns)
+        cols.remove(self.target)
+        cols.remove('date')
+        df_raw = df_raw[['date'] + cols + [self.target]]
+        num_train = int(len(df_raw) * 0.7)
+        num_test = int(len(df_raw) * 0.2)
+        num_vali = len(df_raw) - num_train - num_test
+        border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
+        border2s = [num_train, num_train + num_vali, len(df_raw)]
+        border1 = border1s[self.set_type]
+        border2 = border2s[self.set_type]
+
+        if self.features == 'M' or self.features == 'MS':
+            cols_data = df_raw.columns[1:]
+            df_data = df_raw[cols_data]
+        elif self.features == 'S':
+            df_data = df_raw[[self.target]]
+
+        if self.scale:
+            train_data = df_data[border1s[0]:border2s[0]]
+            self.scaler.fit(train_data.values)
+            data = self.scaler.transform(df_data.values)
+        else:
+            data = df_data.values
+
+        df_stamp = df_raw[['date']][border1:border2]
+        df_stamp['date'] = pd.to_datetime(df_stamp.date)
+        if self.timeenc == 0:
+            df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
+            df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
+            df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
+            df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
+
+
+
+        self.data_x = data[border1:border2]
+        self.data_y = data[border1:border2]
+
+
+    def __getitem__(self, index):
+
+        s_begin = index
+        s_end = s_begin + self.seq_len
+        r_begin = s_end - self.label_len
+        r_end = r_begin + self.label_len + self.pred_len
+
+        seq_x = self.data_x[s_begin:s_end]
+        seq_y = self.data_y[r_begin:r_end]
+        res = []
+        for end in self.seperate:
+            r_end = r_begin + self.label_len + end
+
+            if r_end <= self.data_y.shape[0]:
+                res.append(self.data_y[r_begin:r_end])
+            else:
+                res.append(np.full((r_end - r_begin,self.data_y.shape[-1]), np.nan))
+
+        return seq_x, seq_y, res[0], res[1], res[2], res[3]
+
+    def __len__(self):
+        return len(self.data_x) - self.seq_len - self.raw_start + 1
+
+    def inverse_transform(self, data):
+        return self.scaler.inverse_transform(data)
+
+    
+    
+   

Long_Term_Forecasting/main.py

@@ -0,0 +1,349 @@
+import argparse
+import time
+import datetime
+import os
+import sys
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+import random
+import numpy as np
+
+from  einops import rearrange
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+
+import lightning as L
+from lightning.fabric.strategies import DDPStrategy
+import glob
+from data_loader import Dataset_ETT_hour, Dataset_ETT_minute, Dataset_Custom
+from metrics import metric
+from tqdm import tqdm
+       
+from transformers import AutoModelForCausalLM
+
+
+
+data_dict = {
+    'ETTh1': Dataset_ETT_hour,
+    'ETTh2': Dataset_ETT_hour,
+    'ETTm1': Dataset_ETT_minute,
+    'ETTm2': Dataset_ETT_minute,
+    'custom': Dataset_Custom,
+}
+
+
+
+def data_provider(data, 
+                  root_path,
+                  data_path,
+                  batch_size = 128,
+                  seq_len = 96,
+                  pred_len = 96,
+                  flag= 'test',
+                  dataset = None,
+                  num_workers = 8,
+                  seasonal_patterns = None,
+                  target = 'OT',
+                  features = 'M',
+                 ):
+    
+
+    Data = data_dict[data]
+    shuffle_flag = False
+    drop_last = False
+    batch_size = batch_size 
+    data_set = Data(
+        root_path=root_path,
+        data_path=data_path,
+        flag=flag,
+        size=[seq_len, 0, pred_len],
+        features=features,
+        target=target,
+        timeenc=1,
+        freq='h',
+        seasonal_patterns=seasonal_patterns
+    )
+
+    data_loader = DataLoader(
+        data_set,
+        batch_size=batch_size,
+        shuffle=shuffle_flag,
+        num_workers=num_workers,
+        drop_last=drop_last,
+        pin_memory = True)
+    return data_set, data_loader
+
+datasets_configs = {
+'ETTh1': {'data': 'ETTh1',
+            'root_path': './Long_Term_Forecasting/dataset/ETT-small/',
+            'data_path': 'ETTh1.csv',
+            'batch_size':32,
+            'best_other_mse':'0.390',
+            'best_other_mae':'0.406',
+         },
+'ETTh2': {'data': 'ETTh2',
+            'root_path': './Long_Term_Forecasting/dataset/ETT-small/',
+            'data_path': 'ETTh2.csv',
+            'batch_size':32,
+            'best_other_mse':'0.330',
+            'best_other_mae':'0.375',
+         },
+'ETTm1': {'data': 'ETTm1',
+            'root_path': './Long_Term_Forecasting/dataset/ETT-small/',
+            'data_path': 'ETTm1.csv',
+            'batch_size':32,
+            'best_other_mse':'0.351',
+            'best_other_mae':'0.372',
+         },
+'ETTm2': {'data': 'ETTm2',
+            'root_path': './Long_Term_Forecasting/dataset/ETT-small/',
+            'data_path': 'ETTm2.csv',
+            'batch_size':32,
+            'best_other_mse':'0.255',
+            'best_other_mae':'0.315',
+         },
+'Weather': {'data': 'custom',
+            'root_path': './Long_Term_Forecasting/dataset/weather/',
+            'data_path': 'weather.csv',
+            'batch_size':32,
+            'best_other_mse':'0.226',
+            'best_other_mae':'0.261',
+         }, 
+'Electricity': {'data': 'custom',
+            'root_path': './dataset/',
+            'data_path': 'electricity/electricity.csv',
+            'batch_size':1,
+            'best_other_mse':'0.159',
+            'best_other_mae':'0.253',
+         },
+'Traffic': {'data': 'custom',
+            'root_path': './dataset/',
+            'data_path': 'traffic.csv',
+            'batch_size':32,
+            'best_other_mse':'0.391',
+            'best_other_mae':'0.264',
+         },
+    
+'GlobalTemp': {'data': 'Global_Temp',
+            'root_path': './dataset/',
+            'data_path': 'solar_AL.csv',
+            'batch_size':1,
+            'best_other_mse':'0.322',
+            'best_other_mae':'0.370',
+         },
+
+}
+
+
+if __name__ == '__main__':
+    
+    
+
+
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data', type=str, default='ETTh2', help='dataset type')
+    parser.add_argument('--batch_size', type=int, default=32, help='batch size of train input data')
+    parser.add_argument('--seq_len', type=int, default=96, help='input sequence length')
+    parser.add_argument('--num_gpus', type=int,default=1)
+    parser.add_argument('--future_token', type=int,default=3072)
+    parser.add_argument('-t', '--task_list', action='append')
+    parser.add_argument('--model_name',type=str)
+
+    
+    
+    args = parser.parse_args()
+    
+
+    torch.set_float32_matmul_precision("high")
+
+    
+
+    strategy = DDPStrategy(find_unused_parameters=True)
+    
+    
+    fabric = L.Fabric(devices=args.num_gpus, strategy=strategy)
+    
+    local_rank = fabric.global_rank
+    
+    if local_rank == 0:
+        print(args)
+        
+
+    model = AutoModelForCausalLM.from_pretrained(args.model_name, trust_remote_code=True)
+    model = model.to(local_rank).bfloat16()
+    
+    model = fabric.setup(model)
+    name = args.model_name
+        
+    if local_rank == 0:
+        with open(f'results.txt', 'a') as f:
+                f.write(f"---------------------------------------------------------------------------------")
+    
+  
+    
+    for task_name in args.task_list:
+       
+        task = datasets_configs[task_name]
+        data = task['data']
+        root_path = task['root_path']
+        data_path = task['data_path']
+        best_mse = task['best_other_mse']
+        best_mae = task['best_other_mae']
+        batch_size = min(args.batch_size,task['batch_size'])
+        seq_len = args.seq_len
+        if local_rank == 0:
+            with open(f'results.txt', 'a') as f:
+                seconds_since_epoch = time.time()
+                human_readable_time = datetime.datetime.fromtimestamp(seconds_since_epoch).strftime('%Y-%m-%d %H:%M:%S')
+                f.write(f"{human_readable_time}-------------\n")
+                
+                
+                
+                
+    
+
+        data_set, data_loader = data_provider(data,root_path,data_path,batch_size = batch_size,seq_len=seq_len)
+        data_loader = fabric.setup_dataloaders(data_loader)
+
+
+
+        model.eval()
+        preds = []
+        truths = []
+        preds_s = [[],[],[],[]]
+        truths_s = [[],[],[],[]]
+        intermediates = []
+        xs = []
+
+        seperate_s = [96, 192, 336,720]
+        remains = args.future_token 
+        prevs = 0
+
+        with torch.no_grad():
+            for idx,(x_ori,y,y1,y2,y3,y4) in enumerate(tqdm(data_loader,disable = local_rank != 0)):
+
+                b,c = x_ori.shape[0],x_ori.shape[2]
+                x = rearrange(x_ori, 'b l c -> (b c) l').float().to(local_rank).bfloat16().contiguous()
+                y = rearrange(y1, 'b l c -> (b c) l').float()
+
+                y_s = [y1,y2,y3,y4]
+                res = []
+                res1 = []
+                res2 = []
+                res3 = []
+
+                logits = 0
+                used = 0
+
+
+                for history in [512,1024,2048,4096]:
+                    if history > x.shape[1]:
+                        continue
+                    else:
+                        used += 2
+
+                    x_mean = x[:,-history:].mean(dim = -1,keepdims = True)
+                    x_std = x[:,-history:].std(dim = -1,keepdims = True)
+
+                    x_train = torch.cat((x[:,-history:],-x[:,-history:]),dim = 0)
+
+
+                    logits_all = model(idx = x_train, future_token = args.future_token)
+
+
+                    logits_all = rearrange(logits_all, '(t b) l c d -> b (l c) d t', t = 2)
+                    logits += logits_all[...,0]  -logits_all[...,1].flip(dims = [-1]) 
+
+                logits = logits / used
+
+
+
+                x = torch.cat([x,logits[:,:720,49]],dim = -1).float()
+
+                median = logits[:,:720,49].float()
+                median = median[:,:720] 
+
+                median0 =  rearrange(median, '(b c) l -> b l c',b = b).contiguous().cpu().detach().numpy()
+                y0 = rearrange(y, '(b c) l -> b l c',b = b).contiguous().cpu().detach().numpy()
+
+
+
+                for i, seperate in enumerate(seperate_s):
+                    median_s =  logits[:,:seperate,49].float()
+                    median_s = rearrange(median_s, '(b c) l -> b l c',b = b).contiguous().cpu().detach().numpy()
+                    preds_s[i].append(median_s)
+                    truths_s[i].append(y_s[i].contiguous().cpu().detach().numpy())
+
+
+                xs.append(x_ori.contiguous().cpu().detach().numpy())
+
+
+
+        def gather_losses(loss):
+            """Gather loss values from all GPUs."""
+            if dist.is_initialized():
+                loss_tensor = torch.tensor([loss], device=local_rank)
+                gathered_losses = [torch.zeros_like(loss_tensor) for _ in range(args.num_gpus)]
+                dist.all_gather(gathered_losses, loss_tensor)
+                return torch.cat(gathered_losses).mean()
+            else:
+                return loss
+
+
+
+
+
+
+
+        if local_rank == 0:
+            print(f'Eval on {task_name}-{seq_len}...')
+        mses = []
+        maes = []
+
+
+        for i, seperate in enumerate(seperate_s):
+            if i == 4:
+                break
+            truths = truths_s[i]
+            preds = preds_s[i]
+            truths = np.concatenate(truths,axis = 0)
+            preds = np.concatenate(preds,axis = 0)
+
+            truths = rearrange(truths,'b l c -> b c l')
+            preds = rearrange(preds,'b l c -> b c l')
+            mask = np.isnan(truths).any(axis=2)
+
+            truths1 = truths[~mask]
+            preds1 = preds[~mask]
+
+            truths = rearrange(truths,'b c l-> b l c')
+            preds = rearrange(preds,'b c l-> b l c')
+
+            mae, mse, rmse, mape, mspe = metric(preds1[:,:seperate], truths1[:,:seperate])
+            mae,mse = gather_losses(mae),  gather_losses(mse)
+
+
+            if local_rank == 0:
+                print(f'ours-{name}: mse {mse:.4f} mae {mae:.4f}')
+                mses.append(mse.cpu().numpy())
+                maes.append(mae.cpu().numpy())
+                with open(f'results.txt', 'a') as f:
+
+                    f.write(f"ours-{name}, {data_path.split('.')[0]}-{args.seq_len}-{seperate}-{args.future_token}, mse,  {mse:.5f}, mae, {mae:.5f}\n")
+        if local_rank == 0:
+
+            print(f'ours-{name}-avg: mse {np.mean(mses):.3f} mae {np.mean(maes):.3f}')
+
+        
+            with open(f'results.txt', 'a') as f:
+                f.write(f"ours-{name}, {data_path.split('.')[0]}-avg, mse, {np.mean(mses):.5f}, mae, {np.mean(maes):.5f}\n")
+            print(f'best-avg: mse {best_mse} mae {best_mae}')
+
+
+       
+    dist.destroy_process_group()

Long_Term_Forecasting/metrics.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+
+def RSE(pred, true):
+    return np.sqrt(np.sum((true - pred) ** 2)) / np.sqrt(np.sum((true - true.mean()) ** 2))
+
+
+def CORR(pred, true):
+    u = ((true - true.mean(0)) * (pred - pred.mean(0))).sum(0)
+    d = np.sqrt(((true - true.mean(0)) ** 2 * (pred - pred.mean(0)) ** 2).sum(0))
+    return (u / d).mean(-1)
+
+
+def MAE(pred, true):
+    return np.mean(np.abs(pred - true))
+
+
+def MSE(pred, true):
+    return np.mean((pred - true) ** 2)
+
+
+def RMSE(pred, true):
+    return np.sqrt(MSE(pred, true))
+
+
+def MAPE(pred, true):
+    return np.mean(np.abs((pred - true) / true))
+
+
+def MSPE(pred, true):
+    return np.mean(np.square((pred - true) / true))
+
+
+def metric(pred, true):
+    mae = MAE(pred, true)
+    mse = MSE(pred, true)
+    rmse = RMSE(pred, true)
+    mape = MAPE(pred, true)
+    mspe = MSPE(pred, true)
+
+    return mae, mse, rmse, mape, mspe

README.md

@@ -1,3 +1,9 @@
----
-license: cc-by-4.0
----
+---
+tags:
+- model_hub_mixin
+- pytorch_model_hub_mixin
+---
+
+This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
+- Library: [More Information Needed]
+- Docs: [More Information Needed]

config.json

@@ -0,0 +1,31 @@
+{
+    "architectures": ["YingLong"],
+
+  "auto_map": {
+    "AutoConfig": "model_config.YingLongConfig",
+    "AutoModelForCausalLM": "model.GPT"
+      },
+  "org": "Alibaba",
+  "_mlp_class": "LLaMAMLP",
+  "_norm_class": "FusedRMSNorm",
+  "bias": false,
+  "block_size": 8224,
+  "condense_ratio": 1,
+  "haar_trans": true,
+  "haar_trans_inv": true,
+  "haar_trans_norm": "backward",
+  "intermediate_size": 1024,
+  "n_embd": 256,
+  "n_head": 16,
+  "n_layer": 6,
+  "n_query_groups": 4,
+  "norm_eps": 1e-05,
+  "parallel_residual": false,
+  "patch_size": 32,
+  "quantitle": true,
+  "rope_base": 10000,
+  "rotary_percentage": 1.0,
+  "shared_attention_norm": false,
+  "unet": true,
+  "vocab_size": 1
+}

model.py

@@ -0,0 +1,1526 @@
+"""
+
+Based on the tinyllama implementation: https://github.com/jzhang38/TinyLlama 
+
+"""
+
+
+import math, random
+import numpy as np
+from typing import Any, List, Optional, Tuple
+from typing_extensions import Self
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+from lightning_utilities.core.imports import RequirementCache
+FlashAttention2Available = RequirementCache("flash-attn>=2.0.0.post1")
+
+from flash_attn import flash_attn_func
+from xformers.ops import SwiGLU
+from einops import rearrange
+
+
+from transformers import PreTrainedModel
+from .model_config import YingLongConfig
+
+ 
+    
+    
+class Tokenizer(torch.nn.Module):
+    def __init__(self, config: YingLongConfig, *args,**kwargs) -> None:
+        super().__init__()
+        
+        self.config = config
+        self.tokenizer = nn.Linear(config.patch_size,self.config.n_embd)
+         
+        self.patch_size = config.patch_size
+        self.mask0 = nn.Linear(1,config.n_embd)
+        
+        self.register_buffer('mask_token', torch.zeros(1000))
+        if self.config.haar_trans:
+            
+            self.register_buffer('haar_transform',torch.Tensor(haarMatrix(self.config.patch_size,normalized = self.config.haar_trans_norm)))
+
+        
+        
+    def forward(self,x, 
+                future_token = 0, 
+                prev_token = 0, 
+                factor = 0.2,
+                sequential = False,
+                *args, **kwargs):
+        
+        
+        b = x.shape[0]
+        
+        x_raw = rearrange(x, "b (l c) -> b l c", c = self.patch_size)
+        x_raw_0 = rearrange(x, "b (l c) -> b l c", c = self.patch_size).detach().clone()
+        
+        if future_token == 0:
+            if not sequential:
+                masks = torch.randperm(x_raw.shape[1])
+                unmasks,masks = masks[:int(x_raw.shape[1]*factor)],masks[int(x_raw.shape[1]*factor):]
+            else:
+                masks = [_ for _ in range(x_raw.shape[1])]
+                factor = np.random.rand()*0.6 + 0.2
+                unmasks,masks = masks[:int(x_raw.shape[1]*factor)],masks[int(x_raw.shape[1]*factor):]
+            
+            
+            
+            x_raw_remains = x_raw[:,unmasks,:]
+            
+            mean = x_raw_remains.mean(dim = (-2,-1),keepdims = True)
+            std = x_raw_remains.std(dim = (-2,-1),keepdims = True)
+            x_raw = (x_raw - mean)/ (std + 1e-4)
+          
+            
+            if self.config.haar_trans:
+                x_featured = torch.einsum('blc,ac->bla',x_raw,self.haar_transform)
+                x_featured = self.tokenizer(x_featured)
+            else:
+                x_featured = self.tokenizer(x_raw)
+                
+ 
+            x_featured[:,masks,:] = self.mask0(self.mask_token[0].unsqueeze(0))
+            
+            
+   
+        else:
+   
+            factor = 1
+            more_rows = future_token // self.patch_size + 1
+            prev_more_rows = prev_token // self.patch_size + 1
+        
+            mean = x_raw[:,prev_more_rows:-more_rows,:].mean(dim = (-2,-1),keepdims = True)
+            std = x_raw[:,prev_more_rows:-more_rows,:].std(dim = (-2,-1),keepdims = True)
+            x_raw = (x_raw - mean)/ (std + 1e-4)
+            
+            
+            if self.config.haar_trans:
+                x_featured = torch.einsum('blc,ac->bla',x_raw,self.haar_transform)
+                x_featured = self.tokenizer(x_featured)
+            else:
+                x_featured = self.tokenizer(x_raw)
+    
+    
+            masks = [jj for jj in range(x_featured.shape[1])]
+            masks = masks[-more_rows:]
+    
+            x_featured[:,-more_rows:] = self.mask0(self.mask_token[:len(masks)].unsqueeze(-1)).repeat(x_featured.shape[0],1,1)
+            x_featured[:,:prev_more_rows] = self.mask0(self.mask_token[:prev_more_rows].unsqueeze(-1)).repeat(x_featured.shape[0],1,1)
+
+
+        return x_featured, x_raw_0, masks, mean, std, x_raw
+        
+    
+    
+class model_tmp(PreTrainedModel):
+    config_class = YingLongConfig
+    base_model_prefix = "model"
+    
+    
+    
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=math.sqrt(2.0 / 5 / self.config.n_embd))
+        elif isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=math.sqrt(2.0 / 5 / self.config.n_embd))
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)       
+        for name, p in module.named_parameters():
+            if (name == "proj.weight" and isinstance(module, LLaMAMLP)) or (name == "w3.weight" and isinstance(module, SwiGLU) or (name=="proj.weight" and isinstance(module, BidirectedlSelfAttention))):  
+                nn.init.normal_(p, mean=0.0, std=1 / math.sqrt(self.config.n_embd)  /  self.config.n_layer)
+        
+        
+        
+        
+        
+        
+        
+class GPT(model_tmp):
+    def __init__(self, config: YingLongConfig, *args,**kwargs) -> None:
+        
+        
+        super().__init__(config)
+        
+        self.config = config
+        self.patch_size = config.patch_size
+        self.unet = config.unet
+        
+        
+        if self.config._norm_class == "RMSNorm":
+
+            self.config.norm_class =  RMSNorm
+        elif self.config._norm_class == "FusedRMSNorm":
+            self.config.norm_class =  FusedRMSNorm
+        elif self.config._norm_class == 'BatchNorm':
+            self.config.norm_class =  iBatchNorm
+
+    
+        if self.config._mlp_class == "GptNeoxMLP":
+            self.config.mlp_class =  GptNeoxMLP
+        elif self.config._mlp_class == "LLaMAMLP":
+            self.config.mlp_class =  LLaMAMLP
+        
+                
+
+        
+        self.tokenizer = Tokenizer(config)
+
+                
+        self.lm_head =  nn.Linear(config.n_embd, 99*self.patch_size)
+       
+
+        self.quantitleLoss = quantitleLoss(99,patch_size = self.patch_size)
+        
+        
+        
+        if self.unet:
+            assert config.n_layer%2 == 0
+            self.unet_projection = nn.ModuleList(nn.Sequential(nn.Linear(config.n_embd*2,config.n_embd),
+                                                            config.norm_class(config.n_embd, eps=config.norm_eps),
+                                                            )
+                                            for  _ in range(config.n_layer//2)
+                                           )
+            self.unet_merge = nn.ModuleList(nn.Sequential(nn.Linear(config.n_embd*2,config.n_embd),
+                                                            config.norm_class(config.n_embd, eps=config.norm_eps),
+                                                            )
+                                            for  _ in range(config.n_layer//2)
+                                           )
+        
+       
+        
+        self.transformer = nn.ModuleDict(dict(h = nn.ModuleList(Block(config) 
+                                          for _ in range(config.n_layer))
+                                             )
+                                        )
+        
+       
+    
+        self.rope_cache = None
+
+
+
+    def forward(
+        self, idx: torch.Tensor, 
+        future_token: int = 0,
+        prev_token: int = 0,
+        *args,**kwargs,
+        ) -> torch.Tensor:
+        
+        if future_token > 0:
+            more_rows = future_token // self.patch_size + 1
+            idx = torch.cat((idx,torch.zeros(idx.shape[0],more_rows*self.patch_size).to(idx.device)),dim = -1).bfloat16()
+        if prev_token > 0:
+            more_rows = prev_token // self.patch_size + 1
+            idx = torch.cat((torch.zeros(idx.shape[0],more_rows*self.patch_size).to(idx.device),idx),dim = -1).bfloat16()
+            
+        B, T = idx.size()
+        
+
+
+        block_size = self.config.block_size
+        max_seq_length = T
+            
+        assert max_seq_length <= block_size, f"Cannot attend to {max_seq_length}, block size is only {block_size}"
+        
+
+        self.rope_cache = self.build_rope_cache(idx)
+        cos, sin = self.rope_cache
+
+        cos = cos[:max(T,1024)]
+        sin = sin[:max(T,1024)]
+        
+
+        
+            
+        x,x_raw,masks,mean,std,_ =  self.tokenizer(idx, future_token =future_token,prev_token = prev_token)
+            
+               
+
+        if self.unet:
+            skips = []
+            
+            
+            
+
+        for block_idx in range(len( self.transformer.h)):
+
+
+            block =  self.transformer.h[block_idx]
+
+            if self.unet and block_idx >=len(self.transformer.h) //2:               
+                x = self.unet_projection[block_idx - len(self.transformer.h) //2](torch.cat((skips.pop(),x),dim = -1))
+
+            x = block(x, (cos, sin), max_seq_length)
+
+            if self.unet and block_idx <len(self.transformer.h) //2:
+                skips.append(x)
+                x_delay = torch.cat((x[:,0,:].unsqueeze(1),x[:,:-1,:]),dim = 1)
+                x = self.unet_merge[block_idx](torch.cat((x_delay,x),dim = -1))
+
+     
+                    
+        
+        res = self.lm_head(x)     
+
+        
+            
+        res = rearrange(res,'b c (l1 l2) -> b c l1 l2', l2 = 99)
+        
+
+        
+        if self.config.haar_trans_inv:
+            res = torch.einsum('bcal,ad->bcdl',res,self.tokenizer.haar_transform)
+            if self.config.haar_trans_norm == "backward":
+                res = res / np.sqrt(res.shape[-2])
+            elif self.config.haar_trans_norm == "forward":
+                res = res * np.sqrt(res.shape[-2])
+
+                
+                
+                
+               
+        res = res * (std.unsqueeze(-1) + 1e-4) + mean.unsqueeze(-1)
+        
+        
+            
+              
+        if future_token == 0:
+            return res[:,masks,:,:], x_raw[:,masks,:]
+        else:
+            return res[:,masks,:,:]
+        
+    def generate(self,*args,**kwargs):
+        res = self.forward(*args,**kwargs)
+        res = rearrange(res, 'b l c d -> b (l c) d')
+        return res[:,:kwargs['future_token'],:]
+
+
+    
+    @classmethod
+    def from_name(cls, name: str, **kwargs: Any) -> Self:
+        return cls(Config.from_name(name, **kwargs))
+
+    def build_rope_cache(self, idx: torch.Tensor) :
+        return build_rope_cache(
+            seq_len=self.config.block_size,
+            n_elem=int(self.config.rotary_percentage * self.config.head_size),
+            dtype=torch.bfloat16,
+            device=idx.device,
+            base = self.config.rope_base,
+            condense_ratio=self.config.condense_ratio,
+        )
+
+
+class Block(nn.Module):
+    def __init__(self, config:YingLongConfig) -> None:
+        super().__init__()
+        self.norm_1 = config.norm_class(config.n_embd, eps=config.norm_eps)
+        self.attn = BidirectedlSelfAttention(config)
+        if not config.shared_attention_norm:
+            self.norm_2 = config.norm_class(config.n_embd, eps=config.norm_eps)
+        self.mlp = config.mlp_class(config)
+        self.config = config
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: Optional[Tuple[torch.Tensor, torch.Tensor]],
+        max_seq_length: int,
+        mask: Optional[torch.Tensor] = None,
+        input_pos: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        n_1 = self.norm_1(x)
+        h = self.attn(n_1, rope, max_seq_length, mask, input_pos)
+        if self.config.parallel_residual:
+            n_2 = n_1 if self.config.shared_attention_norm else self.norm_2(x)
+            x = x + h + self.mlp(n_2)
+        else:
+            if self.config.shared_attention_norm:
+                raise NotImplementedError(
+                    "No checkpoint amongst the ones we support uses this configuration"
+                    " (non-parallel residual and shared attention norm)."
+                )
+            
+            x = x + h
+            x = x + self.mlp(self.norm_2(x))
+        return x
+
+
+class BidirectedlSelfAttention(nn.Module):
+    def __init__(self, config:YingLongConfig) -> None:
+        super().__init__()
+        shape = (config.n_head + 2 * config.n_query_groups) * config.head_size
+        self.attn = nn.Linear(config.n_embd, shape, bias=config.bias)
+        self.proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        self.config = config
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: Tuple[torch.Tensor, torch.Tensor],
+        max_seq_length: int,
+        mask: Optional[torch.Tensor] = None,
+        input_pos: Optional[torch.Tensor] = None,
+        ) -> torch.Tensor:
+        
+        
+        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
+
+        qkv = self.attn(x)
+
+        # assemble into a number of query groups to support MHA, MQA and GQA together (see `config.n_query_groups`)
+        q_per_kv = self.config.n_head // self.config.n_query_groups
+        total_qkv = q_per_kv + 2  # each group has 1+ queries, 1 key, and 1 value
+        qkv = qkv.view(B, T, self.config.n_query_groups, total_qkv, self.config.head_size) # (B, T, n_query_groups, total_qkv, hs)
+
+
+        # split batched computation into three
+        q, k, v = qkv.split((q_per_kv, 1, 1), dim=-2)
+
+        q = q.reshape(B,  T, -1, self.config.head_size)  # (B, T, nh_q, hs)
+        k = k.reshape(B,  T, -1, self.config.head_size)  
+        v = v.reshape(B,  T, -1, self.config.head_size)  
+
+        cos, sin = rope
+
+        q = apply_rotary_emb_func(q, cos, sin, False, True)
+        k = apply_rotary_emb_func(k, cos, sin, False, True)
+
+
+        y = self.scaled_dot_product_attention(q, k, v, mask=mask)
+
+        y = y.reshape(B, T, C)  # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.proj(y)
+
+        return y
+    
+    
+
+    def scaled_dot_product_attention(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ):
+        scale = 1.0 / math.sqrt(self.config.head_size)
+        
+        if (
+            FlashAttention2Available
+            and mask is None
+            and q.device.type == "cuda"
+            and q.dtype in (torch.float16, torch.bfloat16)
+        ):
+            from flash_attn import flash_attn_func
+
+            return flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=scale, causal=False)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        if q.size() != k.size():
+            k = k.repeat_interleave(q.shape[1]//k.shape[1], dim=1)
+            v = v.repeat_interleave(q.shape[1]//v.shape[1], dim=1)
+        y = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, dropout_p=0.0, scale=scale, is_causal=False
+        )
+        return y.transpose(1, 2)
+
+
+    
+    
+    
+    
+class quantitleLoss(torch.nn.Module):
+    def __init__(self,  
+                 qSize = 99,
+                 patch_size = 16,
+                 *args,**kwargs):
+        
+        super().__init__()
+        self.qSize = qSize
+        self.patch_size = patch_size
+        
+       
+        q = np.array([i+1 for i in range(self.qSize)])
+        q = q / (self.qSize + 1)
+        q = q.reshape((1,1,-1))
+        
+        q_variance = q*(1-q)
+        
+        self.register_buffer('q', torch.tensor(q))
+        self.register_buffer('q_variance', torch.tensor(q_variance))
+
+        
+    def forward(self, input: torch.Tensor, target: torch.Tensor,rel_loss = False):
+        
+        
+       
+        target = target.unsqueeze(-1)
+        input = input[:,:target.shape[1],:,:]
+        
+        
+        posPart = input - target
+        negPart = -posPart
+        
+        raw_loss = torch.maximum(self.q * negPart, (1-self.q) * posPart)         
+
+        target_absmean = torch.mean(target.abs(),dim = (1,2),keepdims = True)
+        raw_loss = raw_loss  / torch.sqrt(self.q_variance)  / (target_absmean + 1e-4)
+        
+        return torch.mean(raw_loss)
+        
+
+def haarMatrix_unnormalized(n):
+
+    n = 2**np.ceil(np.log2(n))
+    if n > 2:
+        h = haarMatrix(n / 2)
+    else:
+        return np.array([[1, 1], [1, -1]])
+    h_n = np.kron(h, [1, 1])
+    h_i = np.kron(np.eye(len(h)), [1, -1])
+    h = np.vstack((h_n, h_i))
+    return h
+
+def haarMatrix(n,normalized = 'ortho'):
+    h = haarMatrix_unnormalized(n)
+    scaler = np.diag(1/np.sqrt(np.diag(h@h.transpose())))
+    if normalized == 'ortho':
+        return scaler @ h
+    elif normalized == 'forward':
+        return scaler @ h/ np.sqrt(n)
+        
+    else:
+        return scaler @ h *  np.sqrt(n)
+
+
+    
+class GptNeoxMLP(nn.Module):
+    def __init__(self, config:YingLongConfig) -> None:
+        super().__init__()
+        self.fc = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
+        self.proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc(x)
+        x = torch.nn.functional.gelu(x)
+        return self.proj(x)
+
+
+class LLaMAMLP(nn.Module):
+    def __init__(self, config:YingLongConfig) -> None:
+        super().__init__()
+
+        self.swiglu = SwiGLU(config.n_embd,config.intermediate_size, bias=False, _pack_weights=False)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.swiglu(x)
+
+
+def build_rope_cache(
+    seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000, condense_ratio: int = 1
+) -> Tuple[torch.Tensor,torch.Tensor]:
+    """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, device=device) / n_elem))
+
+    # Create position indexes `[0, 1, ..., seq_len - 1]`
+    seq_idx = torch.arange(seq_len, device=device) / condense_ratio
+
+    # Calculate the product of position index and $\theta_i$
+    idx_theta = torch.outer(seq_idx, theta)
+
+    cos, sin = torch.cos(idx_theta), torch.sin(idx_theta)
+
+    # added by peiyuan to ensure same data type with q, k, to use fused rotary embedding
+    if dtype == torch.bfloat16:
+        return cos.bfloat16(), sin.bfloat16()
+    # this is to mimic the behaviour of complex32, else we will get different results
+    if dtype in (torch.float16, torch.bfloat16, torch.int8):
+        return cos.half(), sin.half()
+    return cos, sin
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    head_size = x.size(-1)
+    x1 = x[..., : head_size // 2]  # (B, nh, T, hs/2)
+    x2 = x[..., head_size // 2 :]  # (B, nh, T, hs/2)
+    rotated = torch.cat((-x2, x1), dim=-1)  # (B, nh, T, hs)
+    roped = (x * cos) + (rotated * sin)
+    return roped.type_as(x)
+
+
+
+
+
+
+
+######################################
+#layernorm
+######################################
+
+
+import torch
+# Copyright (c) 2022, Tri Dao.
+# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py AND https://github.com/Dao-AILab/flash-attention/blob/7a983df74215e035e566e37125b0a71e3618f39d/flash_attn/ops/layer_norm.py#L16
+
+import dropout_layer_norm
+import torch
+from torch.nn import init
+
+
+def maybe_align(x, alignment_in_bytes=16):
+    """Assume that x already has last dim divisible by alignment_in_bytes"""
+    # TD [2023-07-04] I'm not 100% sure that clone will align the memory
+    # https://discuss.pytorch.org/t/how-to-ensure-that-tensor-data-ptr-is-aligned-to-16-bytes/183440
+    return x if x.data_ptr() % alignment_in_bytes == 0 else x.clone()
+
+
+def _dropout_add_layer_norm_forward(
+    x0,
+    residual,
+    gamma,
+    beta,
+    rowscale,
+    colscale,
+    dropout_p,
+    epsilon,
+    residual_in_fp32=False,
+    is_rms_norm=False,
+):
+    """Assume that arguments are contiguous and aligned to 16 bytes"""
+    hidden_size = gamma.numel()
+    x0mat = x0.view((-1, hidden_size))
+    residualmat = residual.view((-1, hidden_size)) if residual is not None else None
+    rowscale = rowscale.view(-1) if rowscale is not None else None
+    zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd(
+        x0mat,
+        residualmat,
+        gamma,
+        beta,
+        rowscale,
+        colscale,
+        None,
+        None,
+        dropout_p,
+        epsilon,
+        1.0,
+        0,
+        None,
+        residual_in_fp32,
+        is_rms_norm,
+    )
+    # dmask is None if dropout_p == 0.0
+    # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype
+    return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma
+
+
+def _dropout_add_layer_norm_backward(
+    dz,
+    dx,
+    x,
+    x0,
+    dmask,
+    mu,
+    rsigma,
+    gamma,
+    rowscale,
+    colscale,
+    dropout_p,
+    has_residual,
+    is_rms_norm=False,
+):
+    """Assume that arguments are contiguous and aligned to 16 bytes
+    dx == None means that it was a post-norm architecture
+    (x = drop(x0) + residual was not returned in the fwd).
+    x0 must not be None if we have colscale.
+    """
+    hidden_size = gamma.numel()
+    xmat = x.view((-1, hidden_size))
+    dzmat = dz.view(xmat.shape)
+    dxmat = dx.view(xmat.shape) if dx is not None else None
+    x0mat = x0.view((-1, hidden_size)) if x0 is not None else None
+    rowscale = rowscale.view(-1) if rowscale is not None else None
+    if colscale is not None:
+        assert x0 is not None, "x0 is required to compute the gradient of colscale"
+    dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd(
+        dzmat,
+        dxmat,
+        xmat,
+        x0mat,
+        dmask,
+        mu,
+        rsigma,
+        gamma,
+        rowscale,
+        colscale,
+        None,
+        None,
+        dropout_p,
+        1.0,
+        0,
+        has_residual,
+        is_rms_norm,
+    )
+    # dresidualmat is None if not has_residual
+    if colscale is None:
+        return dx0mat, dresidualmat, dgamma, dbeta
+    else:
+        dcolscale = rest[0]
+        return dx0mat, dresidualmat, dgamma, dbeta, dcolscale
+
+
+def _dropout_add_layer_norm_subset_forward(
+    x0,
+    residual,
+    gamma,
+    beta,
+    colscale,
+    x0_subset,
+    out_subset,
+    dropout_p,
+    epsilon,
+    rowscale_const,
+    out_numrows,
+    residual_in_fp32=False,
+    is_rms_norm=False,
+):
+    """Assume that arguments are contiguous and aligned to 16 bytes"""
+    hidden_size = gamma.numel()
+    x0mat = x0.view((-1, hidden_size))
+    residualmat = residual.view((-1, hidden_size)) if residual is not None else None
+    x0_subset = x0_subset.view(-1) if x0_subset is not None else None
+    out_subset = out_subset.view(-1) if out_subset is not None else None
+    zmat, xmat, dmask, mu, rsigma = dropout_layer_norm.dropout_add_ln_fwd(
+        x0mat,
+        residualmat,
+        gamma,
+        beta,
+        None,
+        colscale,
+        x0_subset,
+        out_subset,
+        dropout_p,
+        epsilon,
+        rowscale_const,
+        out_numrows,
+        None,
+        residual_in_fp32,
+        is_rms_norm,
+    )
+    # dmask is None if dropout_p == 0.0
+    # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype
+    return zmat, xmat if xmat is not None else x0mat, dmask, mu, rsigma
+
+
+def _dropout_add_layer_norm_subset_backward(
+    dz,
+    dx,
+    x,
+    x0,
+    dmask,
+    mu,
+    rsigma,
+    gamma,
+    colscale,
+    x0_subset,
+    out_subset,
+    dropout_p,
+    rowscale_const,
+    x0_numrows,
+    has_residual,
+    is_rms_norm=False,
+):
+    """Assume that arguments are contiguous and aligned to 16 bytes
+    dx == None means that it was a post-norm architecture
+    (x = drop(x0) + residual was not returned in the fwd).
+    x0 must not be None if we have colscale.
+    """
+    hidden_size = gamma.numel()
+    xmat = x.view((-1, hidden_size))
+    dzmat = dz.view(-1, hidden_size)
+    dxmat = dx.view(xmat.shape) if dx is not None else None
+    x0mat = x0.view((-1, hidden_size)) if x0 is not None else None
+    x0_subset = x0_subset.view(-1) if x0_subset is not None else None
+    out_subset = out_subset.view(-1) if out_subset is not None else None
+    if colscale is not None:
+        assert x0 is not None, "x0 is required to compute the gradient of colscale"
+    dx0mat, dresidualmat, dgamma, dbeta, _, _, *rest = dropout_layer_norm.dropout_add_ln_bwd(
+        dzmat,
+        dxmat,
+        xmat,
+        x0mat,
+        dmask,
+        mu,
+        rsigma,
+        gamma,
+        None,
+        colscale,
+        x0_subset,
+        out_subset,
+        dropout_p,
+        rowscale_const,
+        x0_numrows,
+        has_residual,
+        is_rms_norm,
+    )
+    # dresidualmat is None if not has_residual
+    if colscale is None:
+        return dx0mat, dresidualmat, dgamma, dbeta
+    else:
+        dcolscale = rest[0]
+        return dx0mat, dresidualmat, dgamma, dbeta, dcolscale
+
+
+def _dropout_add_layer_norm_parallel_residual_forward(
+    x0,
+    x1,
+    residual,
+    gamma0,
+    beta0,
+    gamma1,
+    beta1,
+    dropout_p,
+    epsilon,
+    residual_in_fp32=False,
+    is_rms_norm=False,
+):
+    """Assume that arguments are contiguous and aligned to 16 bytes"""
+    hidden_size = gamma0.numel()
+    x0mat = x0.view((-1, hidden_size))
+    x1mat = x1.view((-1, hidden_size)) if x1 is not None else None
+    residualmat = residual.view((-1, hidden_size)) if residual is not None else None
+    (
+        z0mat,
+        z1mat,
+        xmat,
+        dmask0,
+        dmask1,
+        mu,
+        rsigma,
+    ) = dropout_layer_norm.dropout_add_ln_parallel_residual_fwd(
+        x0mat,
+        x1mat,
+        residualmat,
+        gamma0,
+        beta0,
+        gamma1,
+        beta1,
+        dropout_p,
+        epsilon,
+        None,
+        residual_in_fp32,
+        is_rms_norm,
+    )
+    # dmask0 and dmask1 are None if dropout_p == 0.0
+    # xmat is None if dropout_p == 0.0 and residual is None and residual_dtype != input_dtype
+    return z0mat, z1mat, xmat if xmat is not None else x0mat, dmask0, dmask1, mu, rsigma
+
+
+def _dropout_add_layer_norm_parallel_residual_backward(
+    dz0,
+    dz1,
+    dx,
+    x,
+    dmask0,
+    dmask1,
+    mu,
+    rsigma,
+    gamma0,
+    gamma1,
+    dropout_p,
+    has_x1,
+    has_residual,
+    is_rms_norm=False,
+):
+    """Assume that arguments are contiguous and aligned to 16 bytes
+    dx == None means that it was a post-norm architecture
+    (x = drop(x0) + residual was not returned in the fwd).
+    """
+    hidden_size = gamma0.numel()
+    xmat = x.view((-1, hidden_size))
+    dz0mat = dz0.view(xmat.shape)
+    dz1mat = dz1.view(xmat.shape) if dz1 is not None else None
+    dxmat = dx.view(xmat.shape) if dx is not None else None
+    (
+        dx0mat,
+        dx1mat,
+        dresidualmat,
+        dgamma0,
+        dbeta0,
+        dgamma1,
+        dbeta1,
+        *rest,
+    ) = dropout_layer_norm.dropout_add_ln_parallel_residual_bwd(
+        dz0mat,
+        dz1mat,
+        dxmat,
+        xmat,
+        dmask0,
+        dmask1,
+        mu,
+        rsigma,
+        gamma0,
+        gamma1,
+        dropout_p,
+        has_x1,
+        has_residual,
+        is_rms_norm,
+    )
+    # dresidualmat is None if not has_residual
+    return dx0mat, dx1mat, dresidualmat, dgamma0, dbeta0, dgamma1, dbeta1
+
+
+class DropoutAddLayerNormFn(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x0,
+        residual,
+        gamma,
+        beta,
+        rowscale,
+        colscale,
+        dropout_p,
+        epsilon,
+        residual_in_fp32=False,
+        prenorm=False,
+        is_rms_norm=False,
+        return_dmask=False,
+    ):
+        x0 = maybe_align(x0.contiguous(), 16)
+        residual = maybe_align(residual.contiguous(), 16) if residual is not None else None
+        gamma = maybe_align(gamma.contiguous(), 16)
+        beta = maybe_align(beta.contiguous(), 16) if beta is not None else None
+        rowscale = maybe_align(rowscale.contiguous(), 16) if rowscale is not None else None
+        colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None
+        zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_forward(
+            x0,
+            residual,
+            gamma,
+            beta,
+            rowscale,
+            colscale,
+            dropout_p,
+            epsilon,
+            residual_in_fp32,
+            is_rms_norm,
+        )
+        # Only need to save x0 if we need to compute gradient wrt colscale
+        x0_saved = x0 if colscale is not None else None
+        ctx.save_for_backward(
+            xmat.view(x0.shape), x0_saved, dmask, gamma, mu, rsigma, rowscale, colscale
+        )
+        ctx.prenorm = prenorm
+        ctx.dropout_p = dropout_p
+        ctx.has_residual = residual is not None
+        ctx.is_rms_norm = is_rms_norm
+        ctx.has_beta = beta is not None
+        if not return_dmask:
+            return (
+                zmat.view(x0.shape) if not prenorm else (zmat.view(x0.shape), xmat.view(x0.shape))
+            )
+        else:
+            dmask = (
+                dmask.view(x0.shape)
+                if dropout_p > 0.0
+                else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device)
+            )
+            ctx.mark_non_differentiable(dmask)
+            return (
+                (zmat.view(x0.shape), dmask)
+                if not prenorm
+                else (zmat.view(x0.shape), xmat.view(x0.shape), dmask)
+            )
+
+    @staticmethod
+    def backward(ctx, dz, *args):
+        # assert dz.is_contiguous()
+        dz = maybe_align(dz.contiguous(), 16)  # this happens!
+        dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None
+        x, x0, dmask, gamma, mu, rsigma, rowscale, colscale = ctx.saved_tensors
+        # x0 is None if colscale is None
+        dropout_p = ctx.dropout_p
+        has_residual = ctx.has_residual
+        dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_backward(
+            dz,
+            dx,
+            x,
+            x0,
+            dmask,
+            mu,
+            rsigma,
+            gamma,
+            rowscale,
+            colscale,
+            dropout_p,
+            has_residual,
+            ctx.is_rms_norm,
+        )
+        dx0 = dx0mat.view(x.shape)
+        dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None
+        dcolscale = rest[0] if colscale is not None else None
+        return (
+            dx0,
+            dresidual,
+            dgamma,
+            dbeta if ctx.has_beta else None,
+            None,
+            dcolscale,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+class DropoutAddLayerNormSubsetFn(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x0,
+        residual,
+        gamma,
+        beta,
+        colscale,
+        x0_subset,
+        out_subset,
+        dropout_p,
+        epsilon,
+        rowscale_const,
+        out_numrows,
+        residual_in_fp32=False,
+        prenorm=False,
+        is_rms_norm=False,
+        return_dmask=False,
+    ):
+        x0 = maybe_align(x0.contiguous(), 16)
+        residual = maybe_align(residual.contiguous(), 16) if residual is not None else None
+        gamma = maybe_align(gamma.contiguous(), 16)
+        beta = maybe_align(beta.contiguous(), 16) if beta is not None else None
+        colscale = maybe_align(colscale.contiguous(), 16) if colscale is not None else None
+        zmat, xmat, dmask, mu, rsigma = _dropout_add_layer_norm_subset_forward(
+            x0,
+            residual,
+            gamma,
+            beta,
+            colscale,
+            x0_subset,
+            out_subset,
+            dropout_p,
+            epsilon,
+            rowscale_const,
+            out_numrows,
+            residual_in_fp32,
+            is_rms_norm,
+        )
+        # Only need to save x0 if we need to compute gradient wrt colscale
+        x0_saved = x0 if colscale is not None else None
+        x_shape = (-1, *x0.shape[1:])
+        ctx.save_for_backward(
+            xmat.view(x_shape), x0_saved, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset
+        )
+        ctx.prenorm = prenorm
+        ctx.dropout_p = dropout_p
+        ctx.rowscale_const = rowscale_const
+        ctx.x0_numrows = x0.shape[:-1].numel()
+        ctx.has_residual = residual is not None
+        ctx.is_rms_norm = is_rms_norm
+        ctx.has_beta = beta is not None
+        z_shape = (-1, *x0.shape[1:])
+        if not return_dmask:
+            return zmat.view(z_shape) if not prenorm else (zmat.view(z_shape), xmat.view(x0.shape))
+        else:
+            z = zmat.view(z_shape)
+            dmask = (
+                dmask.view(x0.shape)
+                if dropout_p > 0.0
+                else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device)
+            )
+            ctx.mark_non_differentiable(dmask)
+            return (z, dmask) if not prenorm else (z, xmat.view(x_shape), dmask)
+
+    @staticmethod
+    def backward(ctx, dz, *args):
+        # assert dz.is_contiguous()
+        dz = maybe_align(dz.contiguous(), 16)  # this happens!
+        dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None
+        x, x0, dmask, gamma, mu, rsigma, colscale, x0_subset, out_subset = ctx.saved_tensors
+        # x0 is None if colscale is None
+        dropout_p = ctx.dropout_p
+        has_residual = ctx.has_residual
+        dx0mat, dresidualmat, dgamma, dbeta, *rest = _dropout_add_layer_norm_subset_backward(
+            dz,
+            dx,
+            x,
+            x0,
+            dmask,
+            mu,
+            rsigma,
+            gamma,
+            colscale,
+            x0_subset,
+            out_subset,
+            dropout_p,
+            ctx.rowscale_const,
+            ctx.x0_numrows,
+            has_residual,
+            ctx.is_rms_norm,
+        )
+        dx0 = dx0mat.view(-1, *x.shape[1:])
+        dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None
+        dcolscale = rest[0] if colscale is not None else None
+        return (
+            dx0,
+            dresidual,
+            dgamma,
+            dbeta if ctx.has_beta else None,
+            dcolscale,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+class DropoutAddLayerNormParallelResidualFn(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x0,
+        x1,
+        residual,
+        gamma0,
+        beta0,
+        gamma1,
+        beta1,
+        dropout_p,
+        epsilon,
+        residual_in_fp32=False,
+        prenorm=False,
+        is_rms_norm=False,
+        return_dmask=False,
+    ):
+        x0 = maybe_align(x0.contiguous(), 16)
+        x1 = maybe_align(x1.contiguous(), 16) if x1 is not None else None
+        residual = maybe_align(residual.contiguous(), 16) if residual is not None else None
+        gamma0 = maybe_align(gamma0.contiguous(), 16)
+        beta0 = maybe_align(beta0.contiguous(), 16) if beta0 is not None else None
+        gamma1 = maybe_align(gamma1.contiguous(), 16) if gamma1 is not None else None
+        beta1 = maybe_align(beta1.contiguous(), 16) if beta1 is not None else None
+        (
+            z0mat,
+            z1mat,
+            xmat,
+            dmask0,
+            dmask1,
+            mu,
+            rsigma,
+        ) = _dropout_add_layer_norm_parallel_residual_forward(
+            x0,
+            x1,
+            residual,
+            gamma0,
+            beta0,
+            gamma1,
+            beta1,
+            dropout_p,
+            epsilon,
+            residual_in_fp32,
+            is_rms_norm,
+        )
+        ctx.save_for_backward(xmat.view(x0.shape), dmask0, dmask1, gamma0, gamma1, mu, rsigma)
+        ctx.prenorm = prenorm
+        ctx.dropout_p = dropout_p
+        ctx.has_x1 = x1 is not None
+        ctx.has_residual = residual is not None
+        ctx.is_rms_norm = is_rms_norm
+        ctx.has_beta = beta0 is not None
+        z = (z0mat.view(x0.shape), z1mat.view(x0.shape) if z1mat is not None else None)
+        if not return_dmask:
+            return z if not prenorm else (*z, xmat.view(x0.shape))
+        else:
+            dmask0 = (
+                dmask0.view(x0.shape)
+                if dropout_p > 0.0
+                else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device)
+            )
+            dmask1 = (
+                dmask1.view(x0.shape)
+                if dropout_p > 0.0 and x1 is not None
+                else torch.ones(x0.shape, dtype=torch.uint8, device=x0.device)
+            )
+            ctx.mark_non_differentiable(dmask0)
+            ctx.mark_non_differentiable(dmask1)
+            return (
+                (*z, dmask0, dmask1) if not prenorm else (*z, xmat.view(x0.shape), dmask0, dmask1)
+            )
+
+    @staticmethod
+    def backward(ctx, dz0, dz1, *args):
+        dz0 = maybe_align(dz0.contiguous(), 16)  # this happens!
+        dz1 = maybe_align(dz1.contiguous(), 16) if dz1 is not None else None
+        dx = maybe_align(args[0].contiguous(), 16) if ctx.prenorm else None
+        x, dmask0, dmask1, gamma0, gamma1, mu, rsigma = ctx.saved_tensors
+        dropout_p = ctx.dropout_p
+        has_x1 = ctx.has_x1
+        has_residual = ctx.has_residual
+        (
+            dx0mat,
+            dx1mat,
+            dresidualmat,
+            dgamma0,
+            dbeta0,
+            dgamma1,
+            dbeta1,
+        ) = _dropout_add_layer_norm_parallel_residual_backward(
+            dz0,
+            dz1,
+            dx,
+            x,
+            dmask0,
+            dmask1,
+            mu,
+            rsigma,
+            gamma0,
+            gamma1,
+            dropout_p,
+            has_x1,
+            has_residual,
+            ctx.is_rms_norm,
+        )
+        dx0 = dx0mat.view(x.shape)
+        dx1 = dx1mat.view(x.shape) if dx1mat is not None else None
+        dresidual = dresidualmat.view(x.shape) if dresidualmat is not None else None
+        return (
+            dx0,
+            dx1,
+            dresidual,
+            dgamma0,
+            dbeta0 if ctx.has_beta else None,
+            dgamma1,
+            dbeta1 if ctx.has_beta else None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+def layer_norm(x, weight, bias, epsilon):
+    return DropoutAddLayerNormFn.apply(x, None, weight, bias, None, None, 0.0, epsilon, False)
+
+
+def dropout_add_layer_norm(
+    x0,
+    residual,
+    weight,
+    bias,
+    dropout_p,
+    epsilon,
+    rowscale=None,
+    layerscale=None,
+    prenorm=False,
+    residual_in_fp32=False,
+    return_dropout_mask=False,
+):
+    """residual_in_fp32 only has an effect if residual is None.
+    Otherwise residual dtype is residual.dtype.
+    """
+    return DropoutAddLayerNormFn.apply(
+        x0,
+        residual,
+        weight,
+        bias,
+        rowscale,
+        layerscale,
+        dropout_p,
+        epsilon,
+        residual_in_fp32,
+        prenorm,
+        False,
+        return_dropout_mask,
+    )
+
+
+def dropout_add_layer_norm_subset(
+    x0,
+    residual,
+    weight,
+    bias,
+    dropout_p,
+    epsilon,
+    layerscale=None,
+    x0_subset=None,
+    out_subset=None,
+    rowscale_const=1.0,
+    out_numrows=0,
+    prenorm=False,
+    residual_in_fp32=False,
+    return_dropout_mask=False,
+):
+    """residual_in_fp32 only has an effect if residual is None.
+    Otherwise residual dtype is residual.dtype.
+    """
+    return DropoutAddLayerNormSubsetFn.apply(
+        x0,
+        residual,
+        weight,
+        bias,
+        layerscale,
+        x0_subset,
+        out_subset,
+        dropout_p,
+        epsilon,
+        rowscale_const,
+        out_numrows,
+        residual_in_fp32,
+        prenorm,
+        False,
+        return_dropout_mask,
+    )
+
+
+def dropout_add_layer_norm_parallel_residual(
+    x0,
+    x1,
+    residual,
+    weight0,
+    bias0,
+    weight1,
+    bias1,
+    dropout_p,
+    epsilon,
+    prenorm=False,
+    residual_in_fp32=False,
+    return_dropout_mask=False,
+):
+    """residual_in_fp32 only has an effect if residual is None.
+    Otherwise residual dtype is residual.dtype.
+    """
+    return DropoutAddLayerNormParallelResidualFn.apply(
+        x0,
+        x1,
+        residual,
+        weight0,
+        bias0,
+        weight1,
+        bias1,
+        dropout_p,
+        epsilon,
+        residual_in_fp32,
+        prenorm,
+        False,
+        return_dropout_mask,
+    )
+
+
+class DropoutAddLayerNorm(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        prenorm=False,
+        p=0.0,
+        eps=1e-5,
+        residual_in_fp32=False,
+        device=None,
+        dtype=None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.prenorm = prenorm
+        self.p = p
+        self.eps = eps
+        self.residual_in_fp32 = residual_in_fp32
+        self.weight = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
+        self.bias = torch.nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
+    def forward(self, x0, residual=None):
+        return dropout_add_layer_norm(
+            x0,
+            residual,
+            self.weight,
+            self.bias,
+            self.p if self.training else 0.0,
+            self.eps,
+            prenorm=self.prenorm,
+            residual_in_fp32=self.residual_in_fp32,
+        )
+        
+def rms_norm(x, weight, epsilon):
+    return DropoutAddLayerNormFn.apply(
+        x, None, weight, None, None, None, 0.0, epsilon, False, False, True
+    )
+class FusedRMSNorm(torch.nn.Module):
+    def __init__(self, size: int, dim: int = -1, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = torch.nn.Parameter(torch.ones(size))
+        self.dim = dim
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init.ones_(self.weight)
+
+    def forward(self, x):
+        return rms_norm(x, self.weight, self.eps)
+    
+    
+class RMSNorm(torch.nn.Module):
+    """Root Mean Square Layer Normalization.
+
+    Derived from https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py. BSD 3-Clause License:
+    https://github.com/bzhangGo/rmsnorm/blob/master/LICENSE.
+    """
+
+    def __init__(self, size: int, dim: int = -1, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(size))
+        self.eps = eps
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # NOTE: the original RMSNorm paper implementation is not equivalent
+        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
+        x_normed = x * torch.rsqrt(norm_x + self.eps)
+        return self.weight * x_normed
+
+    def reset_parameters(self):
+        torch.nn.init.ones_(self.weight)
+
+        
+        
+     
+    
+    
+
+######################################
+#rope_emb
+######################################
+
+
+
+
+   
+        
+        
+# Copyright (c) 2023, Tri Dao.
+
+import math
+from typing import Optional, Tuple
+
+import rotary_emb
+import torch
+from einops import rearrange, repeat
+
+class ApplyRotaryEmb(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, cos, sin, interleaved=False, inplace=False,future_token = 0):
+        """
+            x: (batch_size, seqlen, nheads, headdim)
+            cos, sin: (seqlen, rotary_dim / 2)
+            interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
+                of 1st half and 2nd half (GPT-NeoX style).
+        rotary_dim must be <= headdim
+        Apply rotary embedding to the first rotary_dim of x.
+        """
+        batch, seqlen, nheads, headdim = x.shape
+        rotary_seqlen, rotary_dim = cos.shape
+        rotary_dim *= 2
+
+
+        # print('谁纸盘仲裁',x.shape,cos.shape)
+        # 谁纸盘仲裁 torch.Size([224, 96, 12, 64]) torch.Size([1, 32])
+        # 谁纸盘仲裁 2049 2048
+        assert rotary_dim <= headdim
+        # print(seqlen,rotary_seqlen)
+        assert seqlen <= rotary_seqlen
+        assert sin.shape == (rotary_seqlen, rotary_dim // 2)
+        x_ro = x[..., :rotary_dim]
+        x1, x2 = x_ro.chunk(2, dim=-1) if not interleaved else (x_ro[..., ::2], x_ro[..., 1::2])
+        out = torch.empty_like(x) if not inplace else x
+        out_ro = out[..., :rotary_dim]
+        if inplace:
+            o1, o2 = x1, x2
+        else:
+            o1, o2 = (
+                out_ro.chunk(2, dim=-1)
+                if not interleaved
+                else (out_ro[..., ::2], out_ro[..., 1::2])
+            )
+        rotary_emb.apply_rotary(
+            x1,
+            x2,
+            rearrange(cos[:seqlen], "s d -> s 1 d"),
+            rearrange(sin[:seqlen], "s d -> s 1 d"),
+            o1,
+            o2,
+            False,
+        )
+        if not inplace and rotary_dim < headdim:
+            out[..., rotary_dim:].copy_(x[..., rotary_dim:])
+        ctx.save_for_backward(cos, sin)
+        ctx.interleaved = interleaved
+        ctx.inplace = inplace
+        return out if not inplace else x
+
+    @staticmethod
+    def backward(ctx, do):
+        cos, sin = ctx.saved_tensors
+        _, seqlen, _, headdim = do.shape
+        rotary_dim = cos.shape[-1]
+        rotary_dim *= 2
+        inplace = ctx.inplace
+        do_ro = do[..., :rotary_dim]
+        do1, do2 = (
+            do_ro.chunk(2, dim=-1) if not ctx.interleaved else (do_ro[..., ::2], do_ro[..., 1::2])
+        )
+        dx = torch.empty_like(do) if not inplace else do
+        if inplace:
+            dx1, dx2 = do1, do2
+        else:
+            dx_ro = dx[..., :rotary_dim]
+            dx1, dx2 = (
+                dx_ro.chunk(2, dim=-1)
+                if not ctx.interleaved
+                else (dx_ro[..., ::2], dx_ro[..., 1::2])
+            )
+        rotary_emb.apply_rotary(
+            do1,
+            do2,
+            rearrange(cos[:seqlen], "s d -> s 1 d"),
+            rearrange(sin[:seqlen], "s d -> s 1 d"),
+            dx1,
+            dx2,
+            True,
+        )
+        if not inplace and rotary_dim < headdim:
+            dx[..., rotary_dim:].copy_(do[..., rotary_dim:])
+        return dx, None, None, None, None
+
+
+apply_rotary_emb_func = ApplyRotaryEmb.apply
+

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2df65a7dc7730709b3e1b943fec891d72ef003aa42d0761370f5cd7aa7bf440
+size 14646052

model_config.py

@@ -0,0 +1,100 @@
+from typing import List
+from transformers import PretrainedConfig
+
+
+class YingLongConfig(PretrainedConfig):
+    model_type = "yinglong"
+    # keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        # input_token_len: int = 1,
+        # hidden_size: int = 1024,
+        # intermediate_size: int = 2048,
+        # output_token_lens: List[int] = [1, 8, 32, 64],
+        # num_hidden_layers: int = 8,
+        # num_attention_heads: int = 8,
+        # hidden_act: str = "silu",
+        # use_cache: bool = True,
+        # rope_theta: int = 10000,
+        # attention_dropout: float = 0.0,
+        # initializer_range: float = 0.02,
+        # max_position_embeddings: int = 10000,
+        #####
+        bias = False,
+        condense_ratio = 1,
+        haar_trans = True,
+        haar_trans_inv = True,
+        haar_trans_norm = 'backward',
+        half_diff = False,
+        intermediate_size = 1024,
+        n_embd = 256,
+        n_head = 16,
+        n_layer = 6,
+        n_query_groups = 4,
+        norm_eps = 1e-5,
+        org = 'Alibaba',
+        patch_size = 32,
+        rope_base = 10000,
+        rotary_percentage = 1.0,
+        shared_attention_norm = False,
+        unet = True,
+        _mlp_class = "LLaMAMLP",
+        _norm_class="FusedRMSNorm",
+        *args,
+        **kwargs,
+    ):
+        
+        # self.input_token_len = input_token_len
+        # 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.output_token_lens = output_token_lens;
+        # self.use_cache = use_cache
+        # self.rope_theta = rope_theta
+        # self.attention_dropout = attention_dropout
+        # self.initializer_range = initializer_range
+        # self.max_position_embeddings = max_position_embeddings
+        self.org = 'Alibaba'
+        self.patch_size = patch_size
+        self.unet = unet
+        
+        self.n_embd = n_embd
+        self.intermediate_size = intermediate_size
+        self.n_head = n_head
+        self.n_layer = n_layer
+        self.n_query_groups = n_query_groups
+        self.norm_eps = norm_eps
+        self.bias = bias
+        self.shared_attention_norm = shared_attention_norm
+        
+        self.condense_ratio = condense_ratio
+        self.rope_base = rope_base
+        self.rotary_percentage = rotary_percentage
+        
+        self.haar_trans = haar_trans
+        self.haar_trans_inv = haar_trans_inv
+        self.haar_trans_norm = haar_trans_norm
+        self.half_diff = half_diff
+        
+        self._norm_class = _norm_class
+        
+        self._mlp_class = _mlp_class
+        
+        assert self.n_embd % self.n_head == 0
+        assert self.n_head % self.n_query_groups == 0
+
+        self.head_size = self.n_embd //  self.n_head 
+        self.rope_n_elem = int(self.rotary_percentage * self.head_size)
+        self.rope_condense_ratio = self.condense_ratio
+        
+        
+        
+        
+
+
+        super().__init__(
+            **kwargs,
+        )

未命名.ipynb

@@ -0,0 +1,209 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "94220666-124e-461d-b171-cf1f86056555",
+   "metadata": {
+    "ExecutionIndicator": {
+     "show": false
+    },
+    "execution": {
+     "iopub.execute_input": "2025-05-17T02:39:43.905633Z",
+     "iopub.status.busy": "2025-05-17T02:39:43.905106Z",
+     "iopub.status.idle": "2025-05-17T02:39:43.967707Z",
+     "shell.execute_reply": "2025-05-17T02:39:43.967125Z",
+     "shell.execute_reply.started": "2025-05-17T02:39:43.905610Z"
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import AutoModelForCausalLM\n",
+    "model = AutoModelForCausalLM.from_pretrained('./', trust_remote_code=True,torch_dtype=torch.bfloat16).cuda()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "5feec891-588c-4150-a7c5-6d82e7f62c20",
+   "metadata": {
+    "ExecutionIndicator": {
+     "show": false
+    },
+    "execution": {
+     "iopub.execute_input": "2025-05-17T10:07:59.689508Z",
+     "iopub.status.busy": "2025-05-17T10:07:59.689021Z",
+     "iopub.status.idle": "2025-05-17T10:07:59.806489Z",
+     "shell.execute_reply": "2025-05-17T10:07:59.805876Z",
+     "shell.execute_reply.started": "2025-05-17T10:07:59.689488Z"
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/usr/bin/sh: 1: fabric: not found\n"
+     ]
+    }
+   ],
+   "source": [
+    "batch_size, lookback_length = 1, 2880\n",
+    "seqs = torch.randn(batch_size, lookback_length).bfloat16().cuda()\n",
+    "prediction_length = 96\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 45,
+   "id": "f2b2daff-5bdb-4d07-a9f6-e0ebf538dd58",
+   "metadata": {
+    "ExecutionIndicator": {
+     "show": false
+    },
+    "execution": {
+     "iopub.execute_input": "2025-05-16T21:34:39.066356Z",
+     "iopub.status.busy": "2025-05-16T21:34:39.065849Z",
+     "iopub.status.idle": "2025-05-16T21:34:39.068752Z",
+     "shell.execute_reply": "2025-05-16T21:34:39.068269Z",
+     "shell.execute_reply.started": "2025-05-16T21:34:39.066334Z"
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "output = model.generate(seqs, future_token=prediction_length)\n",
+    "output.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "d02706bd-9ec2-4a61-bbaa-b6c478025e42",
+   "metadata": {
+    "ExecutionIndicator": {
+     "show": false
+    },
+    "execution": {
+     "iopub.execute_input": "2025-05-17T02:39:54.043448Z",
+     "iopub.status.busy": "2025-05-17T02:39:54.043193Z",
+     "iopub.status.idle": "2025-05-17T02:39:54.046722Z",
+     "shell.execute_reply": "2025-05-17T02:39:54.046261Z",
+     "shell.execute_reply.started": "2025-05-17T02:39:54.043431Z"
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "48f905ca-37ae-4382-9da5-8fa4cff50531",
+   "metadata": {
+    "ExecutionIndicator": {
+     "show": false
+    },
+    "execution": {
+     "iopub.execute_input": "2025-05-17T02:39:54.717760Z",
+     "iopub.status.busy": "2025-05-17T02:39:54.717228Z",
+     "iopub.status.idle": "2025-05-17T02:39:56.175917Z",
+     "shell.execute_reply": "2025-05-17T02:39:56.175406Z",
+     "shell.execute_reply.started": "2025-05-17T02:39:54.717720Z"
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "b8eca6b8-3ed1-4d87-b8c7-d84e6994acce",
+   "metadata": {
+    "ExecutionIndicator": {
+     "show": false
+    },
+    "execution": {
+     "iopub.execute_input": "2025-05-17T02:39:57.844387Z",
+     "iopub.status.busy": "2025-05-17T02:39:57.843876Z",
+     "iopub.status.idle": "2025-05-17T02:39:57.849493Z",
+     "shell.execute_reply": "2025-05-17T02:39:57.848990Z",
+     "shell.execute_reply.started": "2025-05-17T02:39:57.844364Z"
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([1, 96, 99])"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "66cf66aa-a991-4a6e-891d-fac9227c383a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3a908d32-4abc-49c6-9ff7-9dfc01379eba",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# !fabric run \\\n",
+    "# --accelerator=cuda \\\n",
+    "# --devices=4 \\\n",
+    "# --num-nodes=1 \\\n",
+    "# --main-port=1145 \\\n",
+    "# Long_Term_Forecasting/main.py \\\n",
+    "# --batch_size 32 \\\n",
+    "# --seq_len 4096 \\\n",
+    "# --future_token 4096 \\\n",
+    "# --model_name ./\\\n",
+    "# --num_gpus 4 \\\n",
+    "# -t ETTh1 \\\n",
+    "# -t ETTh2 \\\n",
+    "# -t ETTm1 \\\n",
+    "# -t ETTm2 \\\n",
+    "# -t Weather "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}