Hum-Works/lodestone-base-4096-v1
12143
1# This training script is a duplicate of the Training.ipynb notebook but can be invoked from the terminal2 3import os4print(os.getcwd())5os.environ["PATH"]="/usr/local/cuda-11.7/bin:"+os.getenv("PATH")6 7os.system('pip uninstall -y torch')8os.system('pip uninstall -y einops')9os.system('pip uninstall -y transformers')10os.system('pip uninstall -y sentence_transformers')11os.system('pip uninstall -y datasets')12os.system('pip uninstall -y sagemaker')13os.system('pip uninstall -y smart_open')14os.system('pip uninstall -y pynvml')15 16os.system('pip install -r lodestone-reqs.txt')17 18os.system('pip install -e ./sentence-transformers')19 20os.system('pip uninstall -y triton')21os.system('pip install --no-deps triton==2.0.0.dev20221202')22 23#####24 25from pynvml import *26import math27from sentence_transformers import models, losses28from sentence_transformers import LoggingHandler, SentenceTransformer, util, InputExample29import logging30import os31import json32import torch33import boto334from smart_open import open35import random36import time37import gc38 39os.environ["PATH"]="/usr/local/cuda-11.7/bin:"+os.getenv("PATH")40os.environ["TOKENIZERS_PARALLELISM"] = "false"41 42#####43 44 45def print_gpu_utilization():46 "This helper function outputs the current GPU memory usage."47 nvmlInit()48 handle = nvmlDeviceGetHandleByIndex(0)49 info = nvmlDeviceGetMemoryInfo(handle)50 return f"GPU memory occupied: {info.used/1024**3} GB."51 52#####53 54 55class MultiDatasetDataLoader:56 """57 This custom dataloader class consumes a list of datasets and a batch size and produces batches randomly sampled 58 from the datasets provided where each batch consists of records from a single dataset and datasets are chosen 59 for batches in proportion to their total number of records.60 """61 def __init__(self, datasets, batch_size_pairs, batch_size_triplets=None, dataset_size_temp=-1, allow_swap=True):62 self.allow_swap = allow_swap63 self.batch_size_pairs = batch_size_pairs64 self.batch_size_triplets = batch_size_pairs if batch_size_triplets is None else batch_size_triplets65 66 # Compute dataset weights67 self.dataset_lengths = list(map(len, datasets))68 self.dataset_lengths_sum = sum(self.dataset_lengths)69 70 weights = []71 # if dataset_size_temp > 0: # Scale probability with dataset size72 # for dataset in datasets:73 # prob = len(dataset) / self.dataset_lengths_sum74 # weights.append(max(1, int(math.pow(prob, 1 / dataset_size_temp) * 1000)))75 # else: # Equal weighting of all datasets76 # weights = [100] * len(datasets)77 for dataset in datasets:78 weights.append(len(dataset))79 80 # logging.info("Dataset lengths and weights: {}".format(list(zip(self.dataset_lengths, weights))))81 82 self.dataset_idx = []83 self.dataset_idx_pointer = 084 85 for idx, weight in enumerate(weights):86 self.dataset_idx.extend([idx] * weight)87 random.shuffle(self.dataset_idx)88 89 self.datasets = []90 for dataset in datasets:91 random.shuffle(dataset)92 self.datasets.append({93 'elements': dataset,94 'pointer': 0,95 })96 97 def __iter__(self):98 for _ in range(int(self.__len__())):99 # Select dataset100 if self.dataset_idx_pointer >= len(self.dataset_idx):101 self.dataset_idx_pointer = 0102 random.shuffle(self.dataset_idx)103 104 dataset_idx = self.dataset_idx[self.dataset_idx_pointer]105 self.dataset_idx_pointer += 1106 107 # Select batch from this dataset108 dataset = self.datasets[dataset_idx]109 batch_size = self.batch_size_pairs if len(dataset['elements'][0].texts) == 2 else self.batch_size_triplets110 111 batch = []112 texts_in_batch = set()113 guid_in_batch = set()114 while len(batch) < batch_size:115 example = dataset['elements'][dataset['pointer']]116 117 valid_example = True118 # First check if one of the texts in already in the batch119 for text in example.texts:120 text_norm = text.strip().lower()121 if text_norm in texts_in_batch:122 valid_example = False123 124 texts_in_batch.add(text_norm)125 126 # If the example has a label, check if label is in batch127 if example.guid is not None:128 valid_example = valid_example and example.guid not in guid_in_batch129 guid_in_batch.add(example.guid)130 131 if valid_example:132 if self.allow_swap and random.random() > 0.5:133 example.texts[0], example.texts[1] = example.texts[1], example.texts[0]134 135 batch.append(example)136 137 dataset['pointer'] += 1138 if dataset['pointer'] >= len(dataset['elements']):139 dataset['pointer'] = 0140 random.shuffle(dataset['elements'])141 142 yield self.collate_fn(batch) if self.collate_fn is not None else batch143 144 def __len__(self):145 return int(self.dataset_lengths_sum / self.batch_size_pairs)146 147#####148 149 150# These four classes of custom generators parse the raw data from the files in S3 and format it into InputExamples which can be properly interpreted by a SentenceTransformer model.151 152class RedditTitleBodyDataset:153 def __init__(self, source_uri, max_seq_length):154 self.source_uri = source_uri155 self.s3_client = boto3.client("s3")156 self.max_seq_length = max_seq_length157 158 def __iter__(self):159 while True:160 for json_line in open(self.source_uri, transport_params={"client": self.s3_client}):161 data_line = json.loads(json_line.strip())162 163 if "title" in data_line and "body" in data_line:164 data = {'guid': None, 'texts': [" ".join(data_line['title'].split(" ")[:self.max_seq_length]), " ".join(data_line['body'].split(" ")[:self.max_seq_length])]}165 record = InputExample(guid=data.get('guid', None), texts=data['texts'])166 167 yield record168 169 170class RedditYearDataset:171 def __init__(self, source_uri, max_seq_length):172 self.source_uri = source_uri173 self.s3_client = boto3.client("s3")174 self.max_seq_length = max_seq_length175 176 def __iter__(self):177 while True:178 for json_line in open(self.source_uri, transport_params={"client": self.s3_client}):179 data_line = json.loads(json_line.strip())180 181 if "response" in data_line and "context" in data_line:182 data = {'guid': None, 'texts': [" ".join(data_line['response'].split(" ")[:self.max_seq_length]), " ".join(data_line['context'].split(" ")[:self.max_seq_length])]}183 record = InputExample(guid=data.get('guid', None), texts=data['texts'])184 185 yield record186 187 188class HuggingFaceQueryPosDataset:189 def __init__(self, source_uri, max_seq_length):190 self.source_uri = source_uri191 self.s3_client = boto3.client("s3")192 self.max_seq_length = max_seq_length193 194 def __iter__(self):195 while True:196 for json_line in open(self.source_uri, transport_params={"client": self.s3_client}):197 data_line = json.loads(json_line.strip())198 199 if "query" in data_line and "pos" in data_line:200 for i in range(len(data_line['pos'])):201 data = {'guid': None, 'texts': [" ".join(data_line['query'].split(" ")[:self.max_seq_length]), " ".join(data_line['pos'][i].split(" ")[:self.max_seq_length])]}202 record = InputExample(guid=data.get('guid', None), texts=data['texts'])203 204 yield record205 206 207class Dataset:208 def __init__(self, source_uri, max_seq_length):209 self.source_uri = source_uri210 self.s3_client = boto3.client("s3")211 self.max_seq_length = max_seq_length212 213 def __iter__(self):214 while True:215 for json_line in open(self.source_uri, transport_params={"client": self.s3_client}):216 data_line = json.loads(json_line.strip())217 218 if not isinstance(data_line, dict):219 data = {'guid': None, 'texts': data_line}220 for text_idx in range(len(data['texts'])):221 data['texts'][text_idx] = " ".join(data['texts'][text_idx].split(" ")[:self.max_seq_length])222 record = InputExample(guid=data.get('guid', None), texts=data['texts'])223 else:224 for text_idx in range(len(data_line['texts'])):225 data_line['texts'][text_idx] = " ".join(data_line['texts'][text_idx].split(" ")[:self.max_seq_length])226 record = InputExample(guid=data_line.get('guid', None), texts=data_line['texts'])227 228 yield record229 230#####231 232 233def build_generators(data_records, max_seq_length=512, testing=False):234 """235 This function consumes the data_records dictionary and creates a new dictionary of data generators where each entry is 236 of the form {filename: data generator object}.237 """238 if testing:239 # filepaths = [file for file in list(data_records.keys()) if file.startswith('S2ORC') or file.startswith('reddit_')]240 filepaths = [file for file in list(data_records.keys())][:3]241 else:242 filepaths = list(data_records.keys())243 generators = {}244 for filepath in filepaths:245 filepath = filepath.strip()246 source_uri = 's3://lodestone-rnd/data/'+filepath247 if filepath in ['S2ORC_citations_abstracts.json.gz', 'amazon-qa.json.gz'] or 'reddit' in filepath:248 if "title" in filepath:249 generators[f'{filepath.split(".")[0]}'] = iter(RedditTitleBodyDataset(source_uri, max_seq_length))250 elif "reddit" in filepath:251 generators[f'{filepath.split(".")[0]}'] = iter(RedditYearDataset(source_uri, max_seq_length))252 else:253 generators[f'{filepath.split(".")[0]}'] = iter(HuggingFaceQueryPosDataset(source_uri, max_seq_length))254 else:255 generators[f'{filepath.split(".")[0]}'] = iter(Dataset(source_uri, max_seq_length))256 257 return generators258 259#####260 261 262def produce_data(data_records, num_chunks, generators, batch_size, failed_on=None, first_iter=False, testing=False, temp=-1):263 """264 This function consumes the data_records dictionary, the number of chunks to break the datasets into, the dictionary of 265 data generators, and a batch size and returns a MultiDatasetDataloader which can be fed into the .fit method of a 266 SentenceTransformer model.267 """268 if testing:269 # filepaths = [file for file in list(data_records.keys()) if file.startswith('S2ORC') or file.startswith('reddit_')]270 filepaths = [file for file in list(data_records.keys())][:3]271 else:272 filepaths = list(data_records.keys())273 datasets = []274 for file_idx, filepath in enumerate(filepaths):275 filepath = filepath.strip()276 dataset = []277 278 if failed_on is not None and failed_on != 1 and first_iter:279 for k in range((failed_on-1)*max(1, data_records[filepath]//num_chunks)):280 next(generators[f'{filepath.split(".")[0]}'])281 for m in range(max(1, data_records[filepath]//num_chunks)):282 dataset.append(next(generators[f'{filepath.split(".")[0]}']))283 else:284 for n in range(max(1, data_records[filepath]//num_chunks)):285 dataset.append(next(generators[f'{filepath.split(".")[0]}']))286 287 datasets.append(dataset)288 logging.info("{}. {}: {}".format(file_idx+1, filepath, len(dataset)))289 290 dataset_lengths_sum = sum(list(map(len, datasets)))291 292 batch_size_pairs = batch_size_triplets = batch_size293 # Special data loader to load from multiple datasets294 train_dataloader = MultiDatasetDataLoader(datasets=datasets,295 batch_size_pairs=batch_size_pairs,296 batch_size_triplets=batch_size_triplets,297 dataset_size_temp=temp)298 299 return train_dataloader, dataset_lengths_sum300 301#####302 303 304def construct_model(model_name, max_seq_length=512):305 """306 This function constructs a SentenceTransformer model from a HuggingFace transformer model name 307 or from a local path to a transformer model repository.308 """309 word_embedding_model = models.Transformer(model_name_or_path=model_name,310 max_seq_length=max_seq_length,311 tokenizer_name_or_path='bert-base-uncased',312 trust_remote_code=True,313 model_args={'torch_dtype': torch.bfloat16})314 pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())315 norm = models.Normalize()316 model = SentenceTransformer(modules=[word_embedding_model, pooling_model, norm], device='cuda')317 model[0].tokenizer.model_max_length = max_seq_length318 319 return model320 321#####322 323 324# Just some code to print debug information to stdout325logging.basicConfig(format='%(asctime)s - %(message)s',326 datefmt='%Y-%m-%d %H:%M:%S',327 level=logging.INFO,328 handlers=[LoggingHandler()])329# /print debug information to stdout330 331#####332 333 334# Set Hyperparameters335model_name = 'mosaic-bert-base-seqlen-2048'336# model_name = 'hum-lodestone-v1'337batch_size = 16338batch_size_pairs = batch_size_triplets = batch_size339max_seq_length = 2048340use_amp = False341 342num_cycles = 2343num_chunks = 50344num_epochs = 2345steps_per_epoch = 10000346# Total training steps = num_cycles * num_chunks * num_epochs * steps_per_epoch = 2 * 50 * 2 * 10,000 = 2,000,000 steps347warmup_steps = 500348 349testing = False350temp = -1351 352#####353 354 355output_path = 'hum-lodestone-v1'356logging.info("Output: "+output_path)357 358# Instantiate SentenceTransformer Model359model = construct_model(model_name=model_name, max_seq_length=max_seq_length)360 361# Load File Names and Record Volumes362with open('data_records.json') as fIn:363 data_records = json.load(fIn)364 365total_pairs = sum(data_records.values())366 367logging.info("Total Training Pairs: {}".format(total_pairs))368 369# Initialize Data Generators370generators = build_generators(data_records=data_records,371 max_seq_length=max_seq_length,372 testing=testing)373 374logging.info("Data Generators Initialized")375 376# Define Training Loss Function377train_loss = losses.MultipleNegativesRankingLoss(model,378 scale=20,379 similarity_fct=util.dot_score)380 381logging.info(print_gpu_utilization())382 383#####384 385 386# Configure Training Cycles387failed_on = None # chunk that the process failed on388random.seed(42)389steps = 0390first_iter = True391for cycle_num in range(num_cycles):392 logging.info("Starting Cycle {}".format(cycle_num+1))393 for chunk_num in range(num_chunks):394 if failed_on is not None and (chunk_num+1) < failed_on and (cycle_num+1) == 1:395 pass396 else:397 logging.info("Chunk {}/{}".format(chunk_num+1, num_chunks))398 logging.info("Loading {} Datasets".format(len([file for file in list(data_records.keys()) if file.startswith('S2ORC') or file.startswith('reddit_')]) if testing else len(data_records)))399 # t_dataload0 = time.time()400 # Create the training dataloader for the given chunk of data401 train_dataloader, dataset_lengths_sum = produce_data(data_records,402 num_chunks,403 generators,404 batch_size,405 failed_on=failed_on,406 first_iter=first_iter,407 testing=testing,408 temp=temp)409 first_iter = False410 # t_dataload1 = time.time()411 # print(t_dataload1-t_dataload0)412 413 logging.info(print_gpu_utilization())414 415 # steps_per_epoch = dataset_lengths_sum // batch_size_pairs416 417 for epoch_num in range(num_epochs):418 logging.info("Performing Cycle {}, Chunk {}, Epoch {}".format(cycle_num+1, chunk_num+1, epoch_num+1))419 try:420 # t_fit0 = time.time()421 # Train the model422 model.fit(train_objectives=[(train_dataloader, train_loss)],423 evaluator=None,424 epochs=1,425 warmup_steps=warmup_steps,426 steps_per_epoch=steps_per_epoch,427 use_amp=use_amp,428 output_path=output_path)429 # t_fit1 = time.time()430 # print(t_fit1-t_fit0)431 432 steps += steps_per_epoch433 434 logging.info(print_gpu_utilization())435 logging.info("Succeeded on Cycle {}, Chunk {}, Epoch {}".format(cycle_num+1, chunk_num+1, epoch_num+1))436 logging.info("{} Steps Completed in Total".format(steps))437 438 with open('train_logs.txt', 'a') as log:439 log.write("Succeeded on Cycle {}, Chunk {}, Epoch {}: {} Steps Completed in Total\n".format(cycle_num+1, chunk_num+1, epoch_num+1, steps))440 441 except:442 logging.info("Failed on Cycle {}, Chunk {}, Epoch {}".format(cycle_num+1, chunk_num+1, epoch_num+1))443 444 with open('train_logs.txt', 'a') as log:445 log.write("Failed on Cycle {}, Chunk {}, Epoch {}: {} Steps Completed in Total\n".format(cycle_num+1, chunk_num+1, epoch_num+1, steps))446 447 finally:448 warmup_steps = 0449 450 # Clear GPU/CUDA memory cache between data chunks451 train_dataloader = None452 model = None453 train_loss = None454 455 gc.collect()456 torch.cuda.empty_cache()457 458 # Reload the model and reinitialize the loss function459 model = construct_model(model_name='hum-lodestone-v1', max_seq_length=max_seq_length)460 461 train_loss = losses.MultipleNegativesRankingLoss(model,462 scale=20,463 similarity_fct=util.dot_score)464 465 logging.info(print_gpu_utilization())