muellerzr/performance-debugging

1# coding=utf-82# Copyright 2021 The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16import evaluate17import torch18from datasets import load_dataset19from torch.optim import AdamW20from torch.utils.data import DataLoader21from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup22 23from accelerate import Accelerator, DistributedType24from accelerate.utils import set_seed25 26import transformers27 28transformers.logging.set_verbosity_error()29 30import os31from torch.nn.parallel import DistributedDataParallel32import torch.distributed as torch_distributed33 34 35 36def get_dataloaders(batch_size: int = 16):37    """38    Creates a set of `DataLoader`s for the `glue` dataset,39    using "bert-base-cased" as the tokenizer.40 41    Args:42        accelerator (`Accelerator`):43            An `Accelerator` object44        batch_size (`int`, *optional*):45            The batch size for the train and validation DataLoaders.46    """47    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")48    datasets = load_dataset("glue", "mrpc")49 50    def tokenize_function(examples):51        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)52        return outputs53 54    tokenized_datasets = datasets.map(55        tokenize_function,56        batched=True,57        remove_columns=["idx", "sentence1", "sentence2"],58    )59    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")60 61    def collate_fn(examples):62        return tokenizer.pad(63            examples,64            padding="longest",65            max_length=None,66            pad_to_multiple_of=8,67            return_tensors="pt",68        )69 70    train_dataloader = DataLoader(71        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True72    )73    eval_dataloader = DataLoader(74        tokenized_datasets["validation"],75        shuffle=False,76        collate_fn=collate_fn,77        batch_size=32,78        drop_last=False,79    )80 81    return train_dataloader, eval_dataloader82 83 84def training_function():85    torch_distributed.init_process_group(backend="nccl")86    num_processes = torch_distributed.get_world_size()87    process_index = torch_distributed.get_rank()88    local_process_index = int(os.environ.get("LOCAL_RANK", -1))89    device = torch.device("cuda", local_process_index)90    torch.cuda.set_device(device)91    config = {"lr": 2e-5, "num_epochs": 3, "seed": 42}92    seed = int(config["seed"])93    batch_size = 32 # Check if this needs to be 32?94    config["batch_size"] = batch_size95    metric = evaluate.load("glue", "mrpc")96 97    set_seed(seed, device_specific=False)98    train_dataloader, eval_dataloader = get_dataloaders(batch_size)99    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True).to(device)100    model = DistributedDataParallel(101        model, device_ids=[local_process_index], output_device=local_process_index102    )103 104    optimizer = AdamW(params=model.parameters(), lr=config["lr"])105    lr_scheduler = get_linear_schedule_with_warmup(106        optimizer=optimizer,107        num_warmup_steps=0,108        num_training_steps=(len(train_dataloader) * config["num_epochs"]),109    )110 111    current_step = 0112    for epoch in range(config["num_epochs"]):113        model.train()114        total_loss = 0115        for _, batch in enumerate(train_dataloader):116            batch = batch.to(device)117            outputs = model(**batch)118            loss = outputs.loss119            total_loss += loss.detach().cpu().float()120            current_step += 1121            loss.backward()122            optimizer.step()123            lr_scheduler.step()124            optimizer.zero_grad()125 126        model.eval()127        for step, batch in enumerate(eval_dataloader):128            # We could avoid this line since we set the accelerator with `device_placement=True`.129            batch = batch.to(device)130            with torch.no_grad():131                outputs = model(**batch)132            predictions = outputs.logits.argmax(dim=-1)133            metric.add_batch(134                predictions=predictions,135                references=batch["labels"],136            )137 138        eval_metric = metric.compute()139        if process_index == 0:140            print(141                f"epoch {epoch}: {eval_metric}\n"142                f"train_loss: {total_loss.item()/len(train_dataloader)}"143            )144 145 146def main():147    training_function()148 149 150if __name__ == "__main__":151    main()152