def main(train_function): # ----- Parse local_rank for torch.distributed.launch ----------- parser = argparse.ArgumentParser() parser.add_argument("--local_rank", type=int) local_rank = parser.parse_args().local_rank if local_rank is None: local_rank = 0 # ----- Setup logging ----------- logger = logging.getLogger(__name__) logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", handlers=[logging.StreamHandler(sys.stdout)], ) logger.setLevel(logging.INFO if is_main_process(local_rank) else logging.WARN) # Set the verbosity to info of the Transformers logger (on main process only): if is_main_process(local_rank): transformers.utils.logging.set_verbosity_info() transformers.utils.logging.enable_default_handler() transformers.utils.logging.enable_explicit_format() # ----- Configurable Params ----------- # List of dicts with configuration for each dataset to be loaded # see available datasets in the Hub: https://huggingface.co/datasets. sizes # are of generated dataset, can be an order of magnitude larger after tokenization. # Not all datasets can be concatenated without preprocessing, features must align datasets_args = [ dict(path="wikitext", name="wikitext-2-raw-v1"), # 12.91 MB # dict(path="wikitext", name="wikitext-103-raw-v1"), # 524 MB # dict(path="ptb_text_only"), # 5.7 MB # dict(path="bookcorpus"), # 4.63 GB # dict(path="wikipedia"), # 35.38 GB ] # Training params # note: in V100 bs=8 uses 11/16 of available gpu mem, bs=12 uses 15/16 output_dir = os.path.expanduser("~/nta/results/bert") training_args = TrainingArguments( # Logging output_dir=output_dir, logging_first_step=True, logging_steps=10, # also define eval_steps eval_steps=10, max_steps=30, # num_train_epochs replaced by steps disable_tqdm=True, run_name="debug_run", # used for wandb, not for Ray # hyperparams per_device_train_batch_size=8, per_device_eval_batch_size=8, learning_rate=1e-4, lr_scheduler_type="linear", warmup_steps=500, weight_decay=1e-6, ) # Evaluate refers to evaluating perplexity on trained model in the validation set # doesn't refer to finetuning and evaluating on downstream tasks such as GLUE seed = random.randint(0, 1000000) # Changing the tokenizer will result in re-tokenizing the dataset. # As a reference, BERT tokenization will take ~ 3 hours for a 5GB dataset config_class = BertConfig tokenizer_name = "bert-base-cased" # ----- Seed ----------- set_seed(seed) print(f"Seed to reproduce: {seed}") # ----- Dataset ----------- # Load multiple datasets and concatenate. # using only 'train' and 'validation' sets, could also include 'test' # if no split is defined, load_dataset returns DatasetDict with all available splits train_datasets = [load_dataset(**args, split="train") for args in datasets_args] val_datasets = [load_dataset(**args, split="validation") for args in datasets_args] dataset = DatasetDict() dataset["train"] = concatenate_datasets(train_datasets) dataset["validation"] = concatenate_datasets(val_datasets) def load_and_split_dataset(dataset_args, split_percentage=5): """Alternative: if no validation set available, manuallly split the train set""" dataset = DatasetDict() dataset["train"] = load_dataset( **dataset_args, split=f"train[{split_percentage}%:]" ) dataset["validation"] = load_dataset( **dataset_args, split=f"train[:{split_percentage}%]" ) return dataset # ----- Load Model ----------- # Load model config = config_class() model = AutoModelForMaskedLM.from_config(config) # Load tokenizer # use_fast falls back to tokenizer lib implementation under the hood tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True) model.resize_token_embeddings(len(tokenizer)) # ----- Preprocess dataset ----------- # Only use the text column name when doing language modeling # this feature might have a different name depending on the dataset # might need to change column names prior to concatenating, if that is the case column_names = dataset["train"].column_names text_column_name = "text" if "text" in column_names else column_names[0] # Setting overwrite_cache to False will retokenize the dataset. # do not overwrite cache if using shared cache repository. overwrite_cache = False preprocessing_num_workers = None # We tokenize every text, then concatenate them together before splitting in smaller # parts. We use `return_special_tokens_mask=True` given # DataCollatorForLanguageModeling is more efficient when it # receives the `special_tokens_mask`. def tokenize_function(examples): return tokenizer(examples[text_column_name], return_special_tokens_mask=True) tokenized_dataset = dataset.map( tokenize_function, batched=True, remove_columns=column_names, num_proc=preprocessing_num_workers, load_from_cache_file=not overwrite_cache, ) # Main data processing function that will concatenate all texts from our dataset and # generate chunks of max_seq_length. max_seq_length = tokenizer.model_max_length def group_texts(examples): # Concatenate all texts. concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()} total_length = len(concatenated_examples[list(examples.keys())[0]]) # We drop the small remainder, we could add padding if the model supported it # instead of this drop, you can customize this part to your needs. total_length = (total_length // max_seq_length) * max_seq_length # Split by chunks of max_len. result = { k: [t[i : i + max_seq_length] for i in range(0, total_length, max_seq_length)] for k, t in concatenated_examples.items() } return result # Note that with `batched=True`, this map processes 1,000 texts together, so # group_texts throws away a remainder for each of those groups of 1,000 texts. # You can adjust batch_size here but a higher value will be slower to preprocess. tokenized_dataset = tokenized_dataset.map( group_texts, batched=True, num_proc=preprocessing_num_workers, load_from_cache_file=not overwrite_cache, ) # Data collator # This one will take care of randomly masking the tokens. # Q: what about dynamic masking, used in Roberta? data_collator = DataCollatorForLanguageModeling( tokenizer=tokenizer, mlm_probability=0.15 ) # ----- Setup Trainer ----------- # Initialize Trainer. Similar to Vernon's Experiment class. # dataloader and training loop are contained in Trainer abstraction trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_dataset["train"], eval_dataset=tokenized_dataset["validation"], tokenizer=tokenizer, data_collator=data_collator, ) # ----- Functions to train and evaluate ----------- if train_function == "huggingface": # Tested run_hf(trainer, logger, output_dir, save_model=True, evaluate=True) elif train_function == "ray_single_node": # Tested run_ray_single_instance( trainer, logger, name="bert_test", config=None, num_samples=1, local_dir=os.path.expanduser("~/nta/results/experiments/transformers"), keep_checkpoints_num=1, resources_per_trial={"cpu": 8}, # note: checkpoint arguments cannot be used with a checkpointable function ) elif train_function == "ray_multiple_nodes": # Untested run_ray_distributed( trainer, logger, name="bert_test", config=None, num_samples=1, local_dir=os.path.expanduser("~/nta/results/experiments/transformers"), keep_checkpoints_num=1, queue_trials=True, verbose=2, resources_per_trial={"gpu": 4}, )
set_seed(seed) print(f"Seed to reproduce: {seed}") # ----- Dataset ----------- # Load multiple datasets and concatenate. # using only 'train' and 'validation' sets, could also include 'test' # if no split is defined, load_dataset returns a DatasetDict with all available splits train_datasets = [ load_dataset(**args, split="train") for args in datasets_args ] val_datasets = [ load_dataset(**args, split="validation") for args in datasets_args ] dataset = DatasetDict() dataset["train"] = concatenate_datasets(train_datasets) dataset["validation"] = concatenate_datasets(val_datasets) def load_and_split_dataset(dataset_args, split_percentage=5): """Alternative: if no validation set available, manuallly split the train set""" dataset = DatasetDict() dataset["train"] = load_dataset(**dataset_args, split=f"train[{split_percentage}%:]") dataset["validation"] = load_dataset(**dataset_args, split=f"train[:{split_percentage}%]") return dataset