| """Utility functions to help setup the model, optimizer, distributed compute, etc. |
| """ |
| import typing |
| import logging |
| from pathlib import Path |
| import sys |
|
|
| import torch |
| import torch.distributed as dist |
| from torch.utils.data import DataLoader, RandomSampler, Dataset |
| from torch.utils.data.distributed import DistributedSampler |
| from ..optimization import AdamW |
|
|
| from ..registry import registry |
|
|
| from .utils import get_effective_batch_size |
| from ._sampler import BucketBatchSampler |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def setup_logging(local_rank: int, |
| save_path: typing.Optional[Path] = None, |
| log_level: typing.Union[str, int] = None) -> None: |
| if log_level is None: |
| level = logging.INFO |
| elif isinstance(log_level, str): |
| level = getattr(logging, log_level.upper()) |
| elif isinstance(log_level, int): |
| level = log_level |
|
|
| if local_rank not in (-1, 0): |
| level = max(level, logging.WARN) |
|
|
| root_logger = logging.getLogger() |
| root_logger.setLevel(level) |
|
|
| formatter = logging.Formatter( |
| "%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
| datefmt="%y/%m/%d %H:%M:%S") |
|
|
| if not root_logger.hasHandlers(): |
| console_handler = logging.StreamHandler(sys.stdout) |
| console_handler.setLevel(level) |
| console_handler.setFormatter(formatter) |
| root_logger.addHandler(console_handler) |
|
|
| if save_path is not None: |
| file_handler = logging.FileHandler(save_path / 'log') |
| file_handler.setLevel(level) |
| file_handler.setFormatter(formatter) |
| root_logger.addHandler(file_handler) |
|
|
|
|
| def setup_optimizer(model, |
| learning_rate: float): |
| """Create the AdamW optimizer for the given model with the specified learning rate. Based on |
| creation in the pytorch_transformers repository. |
| |
| Args: |
| model (PreTrainedModel): The model for which to create an optimizer |
| learning_rate (float): Default learning rate to use when creating the optimizer |
| |
| Returns: |
| optimizer (AdamW): An AdamW optimizer |
| |
| """ |
| param_optimizer = list(model.named_parameters()) |
| no_decay = ['bias', 'gamma', 'beta', 'LayerNorm'] |
| optimizer_grouped_parameters = [ |
| { |
| "params": [ |
| p for n, p in param_optimizer if not any(nd in n for nd in no_decay) |
| ], |
| "weight_decay": 0.01, |
| }, |
| { |
| "params": [ |
| p for n, p in param_optimizer if any(nd in n for nd in no_decay) |
| ], |
| "weight_decay": 0.0, |
| }, |
| ] |
|
|
| optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate) |
| return optimizer |
|
|
|
|
| def setup_dataset(task: str, |
| data_dir: typing.Union[str, Path], |
| split: str, |
| tokenizer: str) -> Dataset: |
| task_spec = registry.get_task_spec(task) |
| return task_spec.dataset(data_dir, split, tokenizer) |
|
|
|
|
| def setup_loader(dataset: Dataset, |
| batch_size: int, |
| local_rank: int, |
| n_gpu: int, |
| gradient_accumulation_steps: int, |
| num_workers: int) -> DataLoader: |
| sampler = DistributedSampler(dataset) if local_rank != -1 else RandomSampler(dataset) |
| batch_size = get_effective_batch_size( |
| batch_size, local_rank, n_gpu, gradient_accumulation_steps) * n_gpu |
| |
| batch_sampler = BucketBatchSampler( |
| sampler, batch_size, False, lambda x: len(x[0]), dataset) |
|
|
| loader = DataLoader( |
| dataset, |
| num_workers=num_workers, |
| collate_fn=dataset.collate_fn, |
| batch_sampler=batch_sampler) |
|
|
| return loader |
|
|
|
|
| def setup_distributed(local_rank: int, |
| no_cuda: bool) -> typing.Tuple[torch.device, int, bool]: |
| if local_rank != -1 and not no_cuda: |
| torch.cuda.set_device(local_rank) |
| device: torch.device = torch.device("cuda", local_rank) |
| n_gpu = 1 |
| dist.init_process_group(backend="nccl") |
| elif not torch.cuda.is_available() or no_cuda: |
| device = torch.device("cpu") |
| n_gpu = 1 |
| else: |
| device = torch.device("cuda") |
| n_gpu = torch.cuda.device_count() |
|
|
| is_master = local_rank in (-1, 0) |
|
|
| return device, n_gpu, is_master |
|
|
