Aluode/PerceptionLabPortable
0
1# Copyright 2020-present the HuggingFace Inc. team.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""15Torch utilities for the Trainer class.16"""17 18import copy19import datetime20import io21import json22import math23import os24import re25import sys26import warnings27from collections.abc import Iterator, Mapping28from contextlib import contextmanager29from dataclasses import dataclass, field30from itertools import chain31from logging import StreamHandler32from typing import Any, Optional, Union33 34import numpy as np35import torch36import torch.distributed as dist37from torch import nn38from torch.utils.data import Dataset, IterableDataset, RandomSampler, Sampler39from torch.utils.data.distributed import DistributedSampler40 41from .integrations.deepspeed import is_deepspeed_zero3_enabled42from .tokenization_utils_base import BatchEncoding43from .utils import (44 is_sagemaker_mp_enabled,45 is_torch_available,46 is_torch_xla_available,47 is_training_run_on_sagemaker,48 logging,49)50 51 52if is_training_run_on_sagemaker():53 logging.add_handler(StreamHandler(sys.stdout))54 55if is_torch_xla_available():56 import torch_xla.runtime as xr57 58if is_torch_available():59 from torch.optim.lr_scheduler import LRScheduler60 61 62logger = logging.get_logger(__name__)63 64 65def get_dataloader_sampler(dataloader):66 if hasattr(dataloader, "batch_sampler") and dataloader.batch_sampler is not None:67 return get_dataloader_sampler(dataloader.batch_sampler)68 elif hasattr(dataloader, "sampler"):69 return dataloader.sampler70 71 72def atleast_1d(tensor_or_array: Union[torch.Tensor, np.ndarray]):73 if isinstance(tensor_or_array, torch.Tensor):74 if hasattr(torch, "atleast_1d"):75 tensor_or_array = torch.atleast_1d(tensor_or_array)76 elif tensor_or_array.ndim < 1:77 tensor_or_array = tensor_or_array[None]78 else:79 tensor_or_array = np.atleast_1d(tensor_or_array)80 return tensor_or_array81 82 83def torch_pad_and_concatenate(tensor1, tensor2, padding_index=-100):84 """Concatenates `tensor1` and `tensor2` on first axis, applying padding on the second if necessary."""85 tensor1 = atleast_1d(tensor1)86 tensor2 = atleast_1d(tensor2)87 88 if len(tensor1.shape) == 1 or tensor1.shape[1] == tensor2.shape[1]:89 return torch.cat((tensor1, tensor2), dim=0)90 91 # Let's figure out the new shape92 new_shape = (tensor1.shape[0] + tensor2.shape[0], max(tensor1.shape[1], tensor2.shape[1])) + tensor1.shape[2:]93 94 # Now let's fill the result tensor95 result = tensor1.new_full(new_shape, padding_index)96 result[: tensor1.shape[0], : tensor1.shape[1]] = tensor197 result[tensor1.shape[0] :, : tensor2.shape[1]] = tensor298 return result99 100 101def numpy_pad_and_concatenate(array1, array2, padding_index=-100):102 """Concatenates `array1` and `array2` on first axis, applying padding on the second if necessary."""103 array1 = atleast_1d(array1)104 array2 = atleast_1d(array2)105 106 if len(array1.shape) == 1 or array1.shape[1] == array2.shape[1]:107 return np.concatenate((array1, array2), axis=0)108 109 # Let's figure out the new shape110 new_shape = (array1.shape[0] + array2.shape[0], max(array1.shape[1], array2.shape[1])) + array1.shape[2:]111 112 # Now let's fill the result tensor113 result = np.full_like(array1, padding_index, shape=new_shape)114 result[: array1.shape[0], : array1.shape[1]] = array1115 result[array1.shape[0] :, : array2.shape[1]] = array2116 return result117 118 119def nested_concat(tensors, new_tensors, padding_index=-100):120 """121 Concat the `new_tensors` to `tensors` on the first dim and pad them on the second if needed. Works for tensors or122 nested list/tuples/dict of tensors.123 """124 if not (isinstance(tensors, torch.Tensor) and isinstance(new_tensors, torch.Tensor)):125 assert type(tensors) is type(new_tensors), (126 f"Expected `tensors` and `new_tensors` to have the same type but found {type(tensors)} and {type(new_tensors)}."127 )128 if isinstance(tensors, (list, tuple)):129 return type(tensors)(nested_concat(t, n, padding_index=padding_index) for t, n in zip(tensors, new_tensors))130 elif isinstance(tensors, torch.Tensor):131 return torch_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index)132 elif isinstance(tensors, Mapping):133 return type(tensors)(134 {k: nested_concat(t, new_tensors[k], padding_index=padding_index) for k, t in tensors.items()}135 )136 elif isinstance(tensors, np.ndarray):137 return numpy_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index)138 else:139 raise TypeError(f"Unsupported type for concatenation: got {type(tensors)}")140 141 142def find_batch_size(tensors):143 """144 Find the first dimension of a tensor in a nested list/tuple/dict of tensors.145 """146 if isinstance(tensors, (list, tuple)):147 for t in tensors:148 result = find_batch_size(t)149 if result is not None:150 return result151 elif isinstance(tensors, Mapping):152 for value in tensors.values():153 result = find_batch_size(value)154 if result is not None:155 return result156 elif isinstance(tensors, (torch.Tensor, np.ndarray)):157 return tensors.shape[0] if len(tensors.shape) >= 1 else None158 159 160def nested_numpify(tensors):161 "Numpify `tensors` (even if it's a nested list/tuple/dict of tensors)."162 if isinstance(tensors, (list, tuple)):163 return type(tensors)(nested_numpify(t) for t in tensors)164 if isinstance(tensors, Mapping):165 return type(tensors)({k: nested_numpify(t) for k, t in tensors.items()})166 167 t = tensors.cpu()168 if t.dtype == torch.bfloat16:169 # As of Numpy 1.21.4, NumPy does not support bfloat16 (see170 # https://github.com/numpy/numpy/blob/a47ecdea856986cd60eabbd53265c2ca5916ad5d/doc/source/user/basics.types.rst ).171 # Until Numpy adds bfloat16, we must convert float32.172 t = t.to(torch.float32)173 return t.numpy()174 175 176def nested_detach(tensors):177 "Detach `tensors` (even if it's a nested list/tuple/dict of tensors)."178 if isinstance(tensors, (list, tuple)):179 return type(tensors)(nested_detach(t) for t in tensors)180 elif isinstance(tensors, Mapping):181 return type(tensors)({k: nested_detach(t) for k, t in tensors.items()})182 return tensors.detach() if isinstance(tensors, torch.Tensor) else tensors183 184 185def nested_xla_mesh_reduce(tensors, name):186 if is_torch_xla_available():187 import torch_xla.core.xla_model as xm188 189 if isinstance(tensors, (list, tuple)):190 return type(tensors)(nested_xla_mesh_reduce(t, f"{name}_{i}") for i, t in enumerate(tensors))191 if isinstance(tensors, Mapping):192 return type(tensors)(193 {k: nested_xla_mesh_reduce(t, f"{name}_{i}") for i, (k, t) in enumerate(tensors.items())}194 )195 196 tensors = atleast_1d(tensors)197 return xm.mesh_reduce(name, tensors, torch.cat)198 else:199 raise ImportError("Torch xla must be installed to use `nested_xla_mesh_reduce`")200 201 202def distributed_concat(tensor: Any, num_total_examples: Optional[int] = None) -> Any:203 try:204 if isinstance(tensor, (tuple, list)):205 return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)206 if isinstance(tensor, Mapping):207 return type(tensor)({k: distributed_concat(t, num_total_examples) for k, t in tensor.items()})208 tensor = atleast_1d(tensor).contiguous()209 output_tensors = [tensor.clone() for _ in range(dist.get_world_size())]210 dist.all_gather(output_tensors, tensor)211 concat = torch.cat(output_tensors, dim=0)212 213 # truncate the dummy elements added by SequentialDistributedSampler214 if num_total_examples is not None:215 concat = concat[:num_total_examples]216 return concat217 except AssertionError:218 raise AssertionError("Not currently using distributed training")219 220 221def distributed_broadcast_scalars(222 scalars: list[Union[int, float]],223 num_total_examples: Optional[int] = None,224 device: Optional[torch.device] = torch.device("cuda"),225) -> torch.Tensor:226 try:227 tensorized_scalar = torch.tensor(scalars, device=device)228 output_tensors = [tensorized_scalar.clone() for _ in range(dist.get_world_size())]229 dist.all_gather(output_tensors, tensorized_scalar)230 concat = torch.cat(output_tensors, dim=0)231 232 # truncate the dummy elements added by SequentialDistributedSampler233 if num_total_examples is not None:234 concat = concat[:num_total_examples]235 return concat236 except AssertionError:237 raise AssertionError("Not currently using distributed training")238 239 240def reissue_pt_warnings(caught_warnings):241 # Reissue warnings242 if len(caught_warnings) > 1:243 for w in caught_warnings:244 if w.category is not UserWarning:245 warnings.warn(w.message, w.category)246 247 248@contextmanager249def torch_distributed_zero_first(local_rank: int):250 """251 Decorator to make all processes in distributed training wait for each local_master to do something.252 253 Args:254 local_rank (`int`): The rank of the local process.255 """256 if local_rank not in [-1, 0]:257 dist.barrier()258 yield259 if local_rank == 0:260 dist.barrier()261 262 263class DistributedSamplerWithLoop(DistributedSampler):264 """265 Like a torch.utils.data.distributed.DistributedSampler` but loops at the end back to the beginning of the shuffled266 samples to make each process have a round multiple of batch_size samples.267 268 Args:269 dataset (`torch.utils.data.Dataset`):270 Dataset used for sampling.271 batch_size (`int`):272 The batch size used with this sampler273 kwargs (`dict[str, Any]`, *optional*):274 All other keyword arguments passed to `DistributedSampler`.275 """276 277 def __init__(self, dataset, batch_size, **kwargs):278 super().__init__(dataset, **kwargs)279 self.batch_size = batch_size280 281 def __iter__(self):282 indices = list(super().__iter__())283 remainder = 0 if len(indices) % self.batch_size == 0 else self.batch_size - len(indices) % self.batch_size284 # DistributedSampler already added samples from the beginning to make the number of samples a round multiple285 # of the world size, so we skip those.286 start_remainder = 1 if self.rank < len(self.dataset) % self.num_replicas else 0287 indices += indices[start_remainder : start_remainder + remainder]288 return iter(indices)289 290 291class EvalLoopContainer:292 """293 Container to store intermediate results of evaluation loop.294 295 Args:296 do_nested_concat (`bool`, *optional*, defaults to `True`):297 If set to `True`, each iteration will recursively concatenate a new object containing tensors to298 the existing stored tensors, provided that the structure of the existing object and the new one299 are identical. If set to `False`, all newly added tensors will be stored in a list.300 padding_index (`int`, *optional*, defaults to -100):301 Value used to pad tensors of different shapes when `do_nested_concat=True`.302 """303 304 def __init__(self, do_nested_concat: bool = True, padding_index: int = -100):305 self.do_nested_concat = do_nested_concat306 self.padding_index = padding_index307 self.tensors = None308 self.arrays = None309 310 def add(self, tensors) -> None:311 """Add tensors to the stored objects. If `do_nested_concat=True`, the tensors will be concatenated recursively."""312 if self.tensors is None:313 self.tensors = tensors if self.do_nested_concat else [tensors]314 elif self.do_nested_concat:315 self.tensors = nested_concat(self.tensors, tensors, padding_index=self.padding_index)316 else:317 self.tensors.append(tensors)318 319 def to_cpu_and_numpy(self) -> None:320 """Move tensors in stored objects to CPU and convert them to numpy arrays."""321 322 # Check if we have something to add, if not just return323 if self.tensors is None:324 return325 326 new_arrays = nested_numpify(self.tensors)327 if self.arrays is None:328 self.arrays = new_arrays329 elif self.do_nested_concat:330 self.arrays = nested_concat(self.arrays, new_arrays, padding_index=self.padding_index)331 else:332 self.arrays.extend(new_arrays)333 334 # reset device tensors after adding to cpu335 self.tensors = None336 337 def get_arrays(self):338 """Returns the numpified and moved to CPU stored objects."""339 self.to_cpu_and_numpy()340 return self.arrays341 342 343class SequentialDistributedSampler(Sampler):344 """345 Distributed Sampler that subsamples indices sequentially, making it easier to collate all results at the end.346 347 Even though we only use this sampler for eval and predict (no training), which means that the model params won't348 have to be synced (i.e. will not hang for synchronization even if varied number of forward passes), we still add349 extra samples to the sampler to make it evenly divisible (like in `DistributedSampler`) to make it easy to `gather`350 or `reduce` resulting tensors at the end of the loop.351 """352 353 def __init__(self, dataset, num_replicas=None, rank=None, batch_size=None):354 warnings.warn(355 "SequentialDistributedSampler is deprecated and will be removed in v5 of Transformers.",356 FutureWarning,357 )358 if num_replicas is None:359 if not dist.is_available():360 raise RuntimeError("Requires distributed package to be available")361 num_replicas = dist.get_world_size()362 if rank is None:363 if not dist.is_available():364 raise RuntimeError("Requires distributed package to be available")365 rank = dist.get_rank()366 self.dataset = dataset367 self.num_replicas = num_replicas368 self.rank = rank369 num_samples = len(self.dataset)370 # Add extra samples to make num_samples a multiple of batch_size if passed371 if batch_size is not None:372 self.num_samples = int(math.ceil(num_samples / (batch_size * num_replicas))) * batch_size373 else:374 self.num_samples = int(math.ceil(num_samples / num_replicas))375 self.total_size = self.num_samples * self.num_replicas376 self.batch_size = batch_size377 378 def __iter__(self):379 indices = list(range(len(self.dataset)))380 381 # add extra samples to make it evenly divisible382 indices += indices[: (self.total_size - len(indices))]383 assert len(indices) == self.total_size, (384 f"Indices length {len(indices)} and total size {self.total_size} mismatched"385 )386 387 # subsample388 indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]389 assert len(indices) == self.num_samples, (390 f"Indices length {len(indices)} and sample number {self.num_samples} mismatched"391 )392 393 return iter(indices)394 395 def __len__(self):396 return self.num_samples397 398 399def get_tpu_sampler(dataset: torch.utils.data.Dataset, batch_size: int):400 if xr.world_size() <= 1:401 return RandomSampler(dataset)402 return DistributedSampler(dataset, num_replicas=xr.world_size(), rank=xr.global_ordinal())403 404 405def nested_new_like(arrays, num_samples, padding_index=-100):406 """Create the same nested structure as `arrays` with a first dimension always at `num_samples`."""407 if isinstance(arrays, (list, tuple)):408 return type(arrays)(nested_new_like(x, num_samples) for x in arrays)409 return np.full_like(arrays, padding_index, shape=(num_samples, *arrays.shape[1:]))410 411 412def expand_like(arrays, new_seq_length, padding_index=-100):413 """Expand the `arrays` so that the second dimension grows to `new_seq_length`. Uses `padding_index` for padding."""414 result = np.full_like(arrays, padding_index, shape=(arrays.shape[0], new_seq_length) + arrays.shape[2:])415 result[:, : arrays.shape[1]] = arrays416 return result417 418 419def nested_truncate(tensors, limit):420 "Truncate `tensors` at `limit` (even if it's a nested list/tuple/dict of tensors)."421 if isinstance(tensors, (list, tuple)):422 return type(tensors)(nested_truncate(t, limit) for t in tensors)423 if isinstance(tensors, Mapping):424 return type(tensors)({k: nested_truncate(t, limit) for k, t in tensors.items()})425 426 return tensors[:limit]427 428 429class DistributedTensorGatherer:430 """431 A class responsible for properly gathering tensors (or nested list/tuple of tensors) on the CPU by chunks.432 433 If our dataset has 16 samples with a batch size of 2 on 3 processes and we gather then transfer on CPU at every434 step, our sampler will generate the following indices:435 436 `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 0, 1]`437 438 to get something of size a multiple of 3 (so that each process gets the same dataset length). Then process 0, 1 and439 2 will be responsible of making predictions for the following samples:440 441 - P0: `[0, 1, 2, 3, 4, 5]`442 - P1: `[6, 7, 8, 9, 10, 11]`443 - P2: `[12, 13, 14, 15, 0, 1]`444 445 The first batch treated on each process will be:446 447 - P0: `[0, 1]`448 - P1: `[6, 7]`449 - P2: `[12, 13]`450 451 So if we gather at the end of the first batch, we will get a tensor (nested list/tuple of tensor) corresponding to452 the following indices:453 454 `[0, 1, 6, 7, 12, 13]`455 456 If we directly concatenate our results without taking any precautions, the user will then get the predictions for457 the indices in this order at the end of the prediction loop:458 459 `[0, 1, 6, 7, 12, 13, 2, 3, 8, 9, 14, 15, 4, 5, 10, 11, 0, 1]`460 461 For some reason, that's not going to roll their boat. This class is there to solve that problem.462 463 Args:464 world_size (`int`):465 The number of processes used in the distributed training.466 num_samples (`int`):467 The number of samples in our dataset.468 make_multiple_of (`int`, *optional*):469 If passed, the class assumes the datasets passed to each process are made to be a multiple of this argument470 (by adding samples).471 padding_index (`int`, *optional*, defaults to -100):472 The padding index to use if the arrays don't all have the same sequence length.473 """474 475 def __init__(self, world_size, num_samples, make_multiple_of=None, padding_index=-100):476 warnings.warn(477 "DistributedTensorGatherer is deprecated and will be removed in v5 of Transformers.",478 FutureWarning,479 )480 self.world_size = world_size481 self.num_samples = num_samples482 total_size = world_size if make_multiple_of is None else world_size * make_multiple_of483 self.total_samples = int(np.ceil(num_samples / total_size)) * total_size484 self.process_length = self.total_samples // world_size485 self._storage = None486 self._offsets = None487 self.padding_index = padding_index488 489 def add_arrays(self, arrays):490 """491 Add `arrays` to the internal storage, Will initialize the storage to the full size at the first arrays passed492 so that if we're bound to get an OOM, it happens at the beginning.493 """494 if arrays is None:495 return496 if self._storage is None:497 self._storage = nested_new_like(arrays, self.total_samples, padding_index=self.padding_index)498 self._offsets = list(range(0, self.total_samples, self.process_length))499 500 slice_len, self._storage = self._nested_set_tensors(self._storage, arrays)501 for i in range(self.world_size):502 self._offsets[i] += slice_len503 504 def _nested_set_tensors(self, storage, arrays):505 if isinstance(arrays, (list, tuple)):506 result = [self._nested_set_tensors(x, y) for x, y in zip(storage, arrays)]507 return result[0][0], type(arrays)(r[1] for r in result)508 assert arrays.shape[0] % self.world_size == 0, (509 f"Arrays passed should all have a first dimension multiple of {self.world_size}, found {arrays.shape[0]}."510 )511 512 slice_len = arrays.shape[0] // self.world_size513 for i in range(self.world_size):514 if len(arrays.shape) == 1:515 storage[self._offsets[i] : self._offsets[i] + slice_len] = arrays[i * slice_len : (i + 1) * slice_len]516 else:517 # Expand the array on the fly if needed.518 if len(storage.shape) > 1 and storage.shape[1] < arrays.shape[1]:519 storage = expand_like(storage, arrays.shape[1], padding_index=self.padding_index)520 storage[self._offsets[i] : self._offsets[i] + slice_len, : arrays.shape[1]] = arrays[521 i * slice_len : (i + 1) * slice_len522 ]523 return slice_len, storage524 525 def finalize(self):526 """527 Return the properly gathered arrays and truncate to the number of samples (since the sampler added some extras528 to get each process a dataset of the same length).529 """530 if self._storage is None:531 return532 if self._offsets[0] != self.process_length:533 logger.warning("Not all data has been set. Are you sure you passed all values?")534 return nested_truncate(self._storage, self.num_samples)535 536 537@dataclass538class LabelSmoother:539 """540 Adds label-smoothing on a pre-computed output from a Transformers model.541 542 Args:543 epsilon (`float`, *optional*, defaults to 0.1):544 The label smoothing factor.545 ignore_index (`int`, *optional*, defaults to -100):546 The index in the labels to ignore when computing the loss.547 """548 549 epsilon: float = 0.1550 ignore_index: int = -100551 552 def __call__(self, model_output, labels, shift_labels=False):553 logits = model_output["logits"] if isinstance(model_output, dict) else model_output[0]554 if shift_labels:555 logits = logits[..., :-1, :].contiguous()556 labels = labels[..., 1:].contiguous()557 558 log_probs = -nn.functional.log_softmax(logits, dim=-1)559 if labels.dim() == log_probs.dim() - 1:560 labels = labels.unsqueeze(-1)561 562 padding_mask = labels.eq(self.ignore_index)563 # In case the ignore_index is -100, the gather will fail, so we replace labels by 0. The padding_mask564 # will ignore them in any case.565 labels = torch.clamp(labels, min=0)566 nll_loss = log_probs.gather(dim=-1, index=labels)567 # works for fp16 input tensor too, by internally upcasting it to fp32568 smoothed_loss = log_probs.sum(dim=-1, keepdim=True, dtype=torch.float32)569 570 nll_loss.masked_fill_(padding_mask, 0.0)571 smoothed_loss.masked_fill_(padding_mask, 0.0)572 573 # Take the mean over the label dimensions, then divide by the number of active elements (i.e. not-padded):574 num_active_elements = padding_mask.numel() - padding_mask.long().sum()575 nll_loss = nll_loss.sum() / num_active_elements576 smoothed_loss = smoothed_loss.sum() / (num_active_elements * log_probs.shape[-1])577 return (1 - self.epsilon) * nll_loss + self.epsilon * smoothed_loss578 579 580def get_length_grouped_indices(lengths, batch_size, mega_batch_mult=None, generator=None):581 """582 Return a list of indices so that each slice of `batch_size` consecutive indices correspond to elements of similar583 lengths. To do this, the indices are:584 585 - randomly permuted586 - grouped in mega-batches of size `mega_batch_mult * batch_size`587 - sorted by length in each mega-batch588 589 The result is the concatenation of all mega-batches, with the batch of `batch_size` containing the element of590 maximum length placed first, so that an OOM happens sooner rather than later.591 """592 # Default for mega_batch_mult: 50 or the number to get 4 megabatches, whichever is smaller.593 if mega_batch_mult is None:594 mega_batch_mult = min(len(lengths) // (batch_size * 4), 50)595 # Just in case, for tiny datasets596 if mega_batch_mult == 0:597 mega_batch_mult = 1598 599 # We need to use torch for the random part as a distributed sampler will set the random seed for torch.600 indices = torch.randperm(len(lengths), generator=generator)601 megabatch_size = mega_batch_mult * batch_size602 megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)]603 megabatches = [sorted(megabatch, key=lambda i: lengths[i], reverse=True) for megabatch in megabatches]604 605 # The rest is to get the biggest batch first.606 # Since each megabatch is sorted by descending length, the longest element is the first607 megabatch_maximums = [lengths[megabatch[0]] for megabatch in megabatches]608 max_idx = torch.argmax(torch.tensor(megabatch_maximums)).item()609 # Switch to put the longest element in first position610 megabatches[0][0], megabatches[max_idx][0] = megabatches[max_idx][0], megabatches[0][0]611 612 return [i for megabatch in megabatches for i in megabatch]613 614 615class LengthGroupedSampler(Sampler):616 r"""617 Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while618 keeping a bit of randomness.619 """620 621 def __init__(622 self,623 batch_size: int,624 dataset: Optional[Dataset] = None,625 lengths: Optional[list[int]] = None,626 model_input_name: Optional[str] = None,627 generator=None,628 ):629 if dataset is None and lengths is None:630 raise ValueError("One of dataset and lengths must be provided.")631 632 self.batch_size = batch_size633 if lengths is None:634 model_input_name = model_input_name if model_input_name is not None else "input_ids"635 if not isinstance(dataset[0], (dict, BatchEncoding)) or model_input_name not in dataset[0]:636 raise ValueError(637 "Can only automatically infer lengths for datasets whose items are dictionaries with an "638 f"'{model_input_name}' key."639 )640 lengths = [len(feature[model_input_name]) for feature in dataset]641 elif isinstance(lengths, torch.Tensor):642 logger.info(643 "If lengths is a torch.Tensor, LengthGroupedSampler will be slow. Converting lengths to list[int]..."644 )645 lengths = lengths.tolist()646 647 self.lengths = lengths648 self.generator = generator649 650 def __len__(self):651 return len(self.lengths)652 653 def __iter__(self):654 indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=self.generator)655 return iter(indices)656 657 658class DistributedLengthGroupedSampler(DistributedSampler):659 r"""660 Distributed Sampler that samples indices in a way that groups together features of the dataset of roughly the same661 length while keeping a bit of randomness.662 """663 664 # Copied and adapted from PyTorch DistributedSampler.665 def __init__(666 self,667 batch_size: int,668 dataset: Optional[Dataset] = None,669 num_replicas: Optional[int] = None,670 rank: Optional[int] = None,671 seed: int = 0,672 drop_last: bool = False,673 lengths: Optional[list[int]] = None,674 model_input_name: Optional[str] = None,675 ):676 if dataset is None and lengths is None:677 raise ValueError("One of dataset and lengths must be provided.")678 if num_replicas is None:679 if not dist.is_available():680 raise RuntimeError("Requires distributed package to be available")681 num_replicas = dist.get_world_size()682 if rank is None:683 if not dist.is_available():684 raise RuntimeError("Requires distributed package to be available")685 rank = dist.get_rank()686 687 self.batch_size = batch_size688 self.num_replicas = num_replicas689 self.rank = rank690 self.epoch = 0691 self.drop_last = drop_last692 693 if lengths is None:694 model_input_name = model_input_name if model_input_name is not None else "input_ids"695 if not isinstance(dataset[0], (dict, BatchEncoding)) or model_input_name not in dataset[0]:696 raise ValueError(697 "Can only automatically infer lengths for datasets whose items are dictionaries with an "698 f"'{model_input_name}' key."699 )700 lengths = [len(feature[model_input_name]) for feature in dataset]701 elif isinstance(lengths, torch.Tensor):702 logger.info(703 "If lengths is a torch.Tensor, DistributedLengthGroupedSampler will be slow. Converting lengths to"704 " list[int]..."705 )706 lengths = lengths.tolist()707 708 self.lengths = lengths709 710 # If the dataset length is evenly divisible by # of replicas, then there711 # is no need to drop any data, since the dataset will be split equally.712 if self.drop_last and len(self.lengths) % self.num_replicas != 0:713 # Split to nearest available length that is evenly divisible.714 # This is to ensure each rank receives the same amount of data when715 # using this Sampler.716 self.num_samples = math.ceil((len(self.lengths) - self.num_replicas) / self.num_replicas)717 else:718 self.num_samples = math.ceil(len(self.lengths) / self.num_replicas)719 self.total_size = self.num_samples * self.num_replicas720 self.seed = seed721 722 def __iter__(self) -> Iterator:723 # Deterministically shuffle based on epoch and seed724 g = torch.Generator()725 g.manual_seed(self.seed + self.epoch)726 indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=g)727 728 if not self.drop_last:729 # add extra samples to make it evenly divisible730 indices += indices[: (self.total_size - len(indices))]731 else:732 # remove tail of data to make it evenly divisible733 indices = indices[: self.total_size]734 assert len(indices) == self.total_size735 736 # subsample737 indices = indices[self.rank : self.total_size : self.num_replicas]738 assert len(indices) == self.num_samples739 740 return iter(indices)741 742 743class ShardSampler(Sampler):744 """745 Sampler that shards batches between several processes. Dispatches indices batch by batch: on 2 processes with batch746 size 4, the first two batches are `[0, 1, 2, 3, 4, 5, 6, 7]` and `[8, 9, 10, 11, 12, 13, 14, 15]`, which shard into747 `[0, 1, 2, 3]` and `[8, 9, 10, 11]` for GPU-0 and `[4, 5, 6, 7]` and `[12, 13, 14, 15]` for GPU-1.748 749 The sampler thus yields `[0, 1, 2, 3, 8, 9, 10, 11]` on GPU-0 and `[4, 5, 6, 7, 12, 13, 14, 15]` on GPU-1.750 """751 752 def __init__(753 self,754 dataset: Dataset,755 batch_size: int = 1,756 drop_last: bool = False,757 num_processes: int = 1,758 process_index: int = 0,759 ):760 self.dataset = dataset761 self.batch_size = batch_size762 self.drop_last = drop_last763 self.num_processes = num_processes764 self.process_index = process_index765 766 self.total_batch_size = total_batch_size = batch_size * num_processes767 768 num_batches = len(dataset) // total_batch_size if drop_last else math.ceil(len(dataset) / total_batch_size)769 self.total_num_samples = num_batches * total_batch_size770 771 def __iter__(self):772 indices = list(range(len(self.dataset)))773 774 # Add extra samples to make it evenly divisible. While loop is there in the edge case we have a tiny dataset775 # and it needs to be done several times.776 while len(indices) < self.total_num_samples:777 indices += indices[: (self.total_num_samples - len(indices))]778 779 result = []780 for batch_start in range(self.batch_size * self.process_index, self.total_num_samples, self.total_batch_size):781 result += indices[batch_start : batch_start + self.batch_size]782 783 return iter(result)784 785 def __len__(self):786 # Each shard only sees a fraction of total_num_samples.787 return self.total_num_samples // self.num_processes788 789 790class IterableDatasetShard(IterableDataset):791 """792 Wraps a PyTorch `IterableDataset` to generate samples for one of the processes only. Instances of this class will793 always yield a number of samples that is a round multiple of the actual batch size (which is `batch_size x794 num_processes`). Depending on the value of the `drop_last` attribute, it will either stop the iteration at the795 first batch that would be too small or loop with indices from the beginning.796 797 On two processes with an iterable dataset yielding of `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]` with a batch size of798 2:799 800 - the shard on process 0 will yield `[0, 1, 4, 5, 8, 9]` so will see batches `[0, 1]`, `[4, 5]`, `[8, 9]`801 - the shard on process 1 will yield `[2, 3, 6, 7, 10, 11]` so will see batches `[2, 3]`, `[6, 7]`, `[10, 11]`802 803 <Tip warning={true}>804 805 If your IterableDataset implements some randomization that needs to be applied the same way on all processes806 (for instance, a shuffling), you should use a `torch.Generator` in a `generator` attribute of the `dataset` to807 generate your random numbers and call the [`~trainer_pt_utils.IterableDatasetShard.set_epoch`] method of this808 object. It will set the seed of this `generator` to `seed + epoch` on all processes before starting the809 iteration. Alternatively, you can also implement a `set_epoch()` method in your iterable dataset to deal with810 this.811 812 </Tip>813 814 Args:815 dataset (`torch.utils.data.IterableDataset`):816 The batch sampler to split in several shards.817 batch_size (`int`, *optional*, defaults to 1):818 The size of the batches per shard.819 drop_last (`bool`, *optional*, defaults to `False`):820 Whether or not to drop the last incomplete batch or complete the last batches by using the samples from the821 beginning.822 num_processes (`int`, *optional*, defaults to 1):823 The number of processes running concurrently.824 process_index (`int`, *optional*, defaults to 0):825 The index of the current process.826 seed (`int`, *optional*, defaults to 0):827 A random seed that will be used for the random number generation in828 [`~trainer_pt_utils.IterableDatasetShard.set_epoch`].829 """830 831 def __init__(832 self,833 dataset: IterableDataset,834 batch_size: int = 1,835 drop_last: bool = False,836 num_processes: int = 1,837 process_index: int = 0,838 seed: int = 0,839 ):840 self.dataset = dataset841 self.batch_size = batch_size842 self.drop_last = drop_last843 self.num_processes = num_processes844 self.process_index = process_index845 self.seed = seed846 self.epoch = 0847 self.num_examples = 0848 849 def set_epoch(self, epoch):850 self.epoch = epoch851 if hasattr(self.dataset, "set_epoch"):852 self.dataset.set_epoch(epoch)853 854 def __iter__(self):855 self.num_examples = 0856 if (857 not hasattr(self.dataset, "set_epoch")858 and hasattr(self.dataset, "generator")859 and isinstance(self.dataset.generator, torch.Generator)860 ):861 self.dataset.generator.manual_seed(self.seed + self.epoch)862 real_batch_size = self.batch_size * self.num_processes863 process_slice = range(self.process_index * self.batch_size, (self.process_index + 1) * self.batch_size)864 865 first_batch = None866 current_batch = []867 for element in self.dataset:868 self.num_examples += 1869 current_batch.append(element)870 # Wait to have a full batch before yielding elements.871 if len(current_batch) == real_batch_size:872 for i in process_slice:873 yield current_batch[i]874 if first_batch is None:875 first_batch = current_batch.copy()876 current_batch = []877 878 # Finished if drop_last is True, otherwise complete the last batch with elements from the beginning.879 if not self.drop_last and len(current_batch) > 0:880 if first_batch is None:881 first_batch = current_batch.copy()882 while len(current_batch) < real_batch_size:883 current_batch += first_batch884 for i in process_slice:885 yield current_batch[i]886 887 def __len__(self):888 # Will raise an error if the underlying dataset is not sized.889 if self.drop_last:890 return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size891 else:892 return math.ceil(len(self.dataset) / (self.batch_size * self.num_processes)) * self.batch_size893 894 895# In order to keep `trainer.py` compact and easy to understand, place any secondary PT Trainer896# helper methods here897 898 899def _get_learning_rate(self):900 if self.is_deepspeed_enabled:901 # with deepspeed's fp16 and dynamic loss scale enabled the optimizer/scheduler steps may902 # not run for the first few dozen steps while loss scale is too large, and thus during903 # that time `get_last_lr` will fail if called during that warm up stage, so work around it:904 try:905 last_lr = self.lr_scheduler.get_last_lr()[0]906 except AssertionError as e:907 if "need to call step" in str(e):908 logger.warning("tried to get lr value before scheduler/optimizer started stepping, returning lr=0")909 last_lr = 0910 else:911 raise912 else:913 if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):914 last_lr = self.optimizer.param_groups[0]["lr"]915 else:916 last_lr = self.lr_scheduler.get_last_lr()[0]917 918 if torch.is_tensor(last_lr):919 last_lr = last_lr.item()920 return last_lr921 922 923def _secs2timedelta(secs):924 """925 Convert seconds to hh:mm:ss.msec, msecs rounded to 2 decimal places.926 """927 928 msec = int(abs(secs - int(secs)) * 100)929 return f"{datetime.timedelta(seconds=int(secs))}.{msec:02d}"930 931 932def metrics_format(metrics: dict[str, float]) -> dict[str, float]:933 """934 Reformat Trainer metrics values to a human-readable format.935 936 Args:937 metrics (`dict[str, float]`):938 The metrics returned from train/evaluate/predict939 940 Returns:941 metrics (`dict[str, float]`): The reformatted metrics942 """943 944 metrics_copy = metrics.copy()945 for k, v in metrics_copy.items():946 if "_mem_" in k:947 metrics_copy[k] = f"{v >> 20}MB"948 elif "_runtime" in k:949 metrics_copy[k] = _secs2timedelta(v)950 elif k == "total_flos":951 metrics_copy[k] = f"{int(v) >> 30}GF"952 elif isinstance(metrics_copy[k], float):953 metrics_copy[k] = round(v, 4)954 955 return metrics_copy956 957 958def log_metrics(self, split, metrics):959 """960 Log metrics in a specially formatted way.961 962 Under distributed environment this is done only for a process with rank 0.963 964 Args:965 split (`str`):966 Mode/split name: one of `train`, `eval`, `test`967 metrics (`dict[str, float]`):968 The metrics returned from train/evaluate/predictmetrics: metrics dict969 970 Notes on memory reports:971 972 In order to get memory usage report you need to install `psutil`. You can do that with `pip install psutil`.973 974 Now when this method is run, you will see a report that will include:975 976 ```977 init_mem_cpu_alloc_delta = 1301MB978 init_mem_cpu_peaked_delta = 154MB979 init_mem_gpu_alloc_delta = 230MB980 init_mem_gpu_peaked_delta = 0MB981 train_mem_cpu_alloc_delta = 1345MB982 train_mem_cpu_peaked_delta = 0MB983 train_mem_gpu_alloc_delta = 693MB984 train_mem_gpu_peaked_delta = 7MB985 ```986 987 **Understanding the reports:**988 989 - the first segment, e.g., `train__`, tells you which stage the metrics are for. Reports starting with `init_`990 will be added to the first stage that gets run. So that if only evaluation is run, the memory usage for the991 `__init__` will be reported along with the `eval_` metrics.992 - the third segment, is either `cpu` or `gpu`, tells you whether it's the general RAM or the gpu0 memory993 metric.994 - `*_alloc_delta` - is the difference in the used/allocated memory counter between the end and the start of the995 stage - it can be negative if a function released more memory than it allocated.996 - `*_peaked_delta` - is any extra memory that was consumed and then freed - relative to the current allocated997 memory counter - it is never negative. When you look at the metrics of any stage you add up `alloc_delta` +998 `peaked_delta` and you know how much memory was needed to complete that stage.999 1000 The reporting happens only for process of rank 0 and gpu 0 (if there is a gpu). Typically this is enough since the1001 main process does the bulk of work, but it could be not quite so if model parallel is used and then other GPUs may1002 use a different amount of gpu memory. This is also not the same under DataParallel where gpu0 may require much more1003 memory than the rest since it stores the gradient and optimizer states for all participating GPUs. Perhaps in the1004 future these reports will evolve to measure those too.1005 1006 The CPU RAM metric measures RSS (Resident Set Size) includes both the memory which is unique to the process and the1007 memory shared with other processes. It is important to note that it does not include swapped out memory, so the1008 reports could be imprecise.1009 1010 The CPU peak memory is measured using a sampling thread. Due to python's GIL it may miss some of the peak memory if1011 that thread didn't get a chance to run when the highest memory was used. Therefore this report can be less than1012 reality. Using `tracemalloc` would have reported the exact peak memory, but it doesn't report memory allocations1013 outside of python. So if some C++ CUDA extension allocated its own memory it won't be reported. And therefore it1014 was dropped in favor of the memory sampling approach, which reads the current process memory usage.1015 1016 The GPU allocated and peak memory reporting is done with `torch.cuda.memory_allocated()` and1017 `torch.cuda.max_memory_allocated()`. This metric reports only "deltas" for pytorch-specific allocations, as1018 `torch.cuda` memory management system doesn't track any memory allocated outside of pytorch. For example, the very1019 first cuda call typically loads CUDA kernels, which may take from 0.5 to 2GB of GPU memory.1020 1021 Note that this tracker doesn't account for memory allocations outside of [`Trainer`]'s `__init__`, `train`,1022 `evaluate` and `predict` calls.1023 1024 Because `evaluation` calls may happen during `train`, we can't handle nested invocations because1025 `torch.cuda.max_memory_allocated` is a single counter, so if it gets reset by a nested eval call, `train`'s tracker1026 will report incorrect info. If this [pytorch issue](https://github.com/pytorch/pytorch/issues/16266) gets resolved1027 it will be possible to change this class to be re-entrant. Until then we will only track the outer level of1028 `train`, `evaluate` and `predict` methods. Which means that if `eval` is called during `train`, it's the latter1029 that will account for its memory usage and that of the former.1030 1031 This also means that if any other tool that is used along the [`Trainer`] calls1032 `torch.cuda.reset_peak_memory_stats`, the gpu peak memory stats could be invalid. And the [`Trainer`] will disrupt1033 the normal behavior of any such tools that rely on calling `torch.cuda.reset_peak_memory_stats` themselves.1034 1035 For best performance you may want to consider turning the memory profiling off for production runs.1036 """1037 if not self.is_world_process_zero():1038 return1039 1040 print(f"***** {split} metrics *****")1041 metrics_formatted = metrics_format(metrics)1042 k_width = max(len(str(x)) for x in metrics_formatted)1043 v_width = max(len(str(x)) for x in metrics_formatted.values())1044 for key in sorted(metrics_formatted.keys()):1045 print(f" {key: <{k_width}} = {metrics_formatted[key]:>{v_width}}")1046 1047 1048def save_metrics(self, split, metrics, combined=True):1049 """1050 Save metrics into a json file for that split, e.g. `train_results.json`.1051 1052 Under distributed environment this is done only for a process with rank 0.1053 1054 Args:1055 split (`str`):1056 Mode/split name: one of `train`, `eval`, `test`, `all`1057 metrics (`dict[str, float]`):1058 The metrics returned from train/evaluate/predict1059 combined (`bool`, *optional*, defaults to `True`):1060 Creates combined metrics by updating `all_results.json` with metrics of this call1061 1062 To understand the metrics please read the docstring of [`~Trainer.log_metrics`]. The only difference is that raw1063 unformatted numbers are saved in the current method.1064 1065 """1066 if not self.is_world_process_zero():1067 return1068 1069 path = os.path.join(self.args.output_dir, f"{split}_results.json")1070 with open(path, "w") as f:1071 json.dump(metrics, f, indent=4, sort_keys=True)1072 1073 if combined:1074 path = os.path.join(self.args.output_dir, "all_results.json")1075 if os.path.exists(path):1076 with open(path) as f:1077 all_metrics = json.load(f)1078 else:1079 all_metrics = {}1080 1081 all_metrics.update(metrics)1082 with open(path, "w") as f:1083 json.dump(all_metrics, f, indent=4, sort_keys=True)1084 1085 1086def save_state(self):1087 """1088 Saves the Trainer state, since Trainer.save_model saves only the tokenizer with the model.1089 1090 Under distributed environment this is done only for a process with rank 0.1091 """1092 if not self.is_world_process_zero():1093 return1094 1095 path = os.path.join(self.args.output_dir, "trainer_state.json")1096 self.state.save_to_json(path)1097 1098 1099def get_model_param_count(model, trainable_only=False):1100 """1101 Calculate model's total param count. If trainable_only is True then count only those requiring grads.1102 """1103 if is_deepspeed_zero3_enabled():1104 1105 def numel(p):1106 return p.ds_numel if hasattr(p, "ds_numel") else p.numel()1107 1108 else:1109 1110 def numel(p):1111 return p.numel()1112 1113 return sum(numel(p) for p in model.parameters() if not trainable_only or p.requires_grad)1114 1115 1116def get_parameter_names(model, forbidden_layer_types, forbidden_layer_names=None):1117 """1118 Returns the names of the model parameters that are not inside a forbidden layer.1119 """1120 forbidden_layer_patterns = (1121 [re.compile(pattern) for pattern in forbidden_layer_names] if forbidden_layer_names is not None else []1122 )1123 result = []1124 for name, child in model.named_children():1125 child_params = get_parameter_names(child, forbidden_layer_types, forbidden_layer_names)1126 result += [1127 f"{name}.{n}"1128 for n in child_params1129 if not isinstance(child, tuple(forbidden_layer_types))1130 and not any(pattern.search(f"{name}.{n}".lower()) for pattern in forbidden_layer_patterns)1131 ]1132 # Add model specific parameters that are not in any child1133 result += [1134 k for k in model._parameters if not any(pattern.search(k.lower()) for pattern in forbidden_layer_patterns)1135 ]1136 1137 return result1138 1139 1140def get_module_class_from_name(module, name):1141 """1142 Gets a class from a module by its name.1143 1144 Args:1145 module (`torch.nn.Module`): The module to get the class from.1146 name (`str`): The name of the class.1147 """1148 modules_children = list(module.children())1149 if module.__class__.__name__ == name:1150 return module.__class__1151 elif len(modules_children) == 0:1152 return1153 else:1154 for child_module in modules_children:1155 module_class = get_module_class_from_name(child_module, name)1156 if module_class is not None:1157 return module_class1158 1159 1160def remove_dummy_checkpoint(is_main_process, output_dir, filenames):1161 if is_main_process:1162 for filename in filenames:1163 file = os.path.join(output_dir, filename)1164 if os.path.isfile(file):1165 os.remove(file)1166 1167 1168if is_sagemaker_mp_enabled():1169 import smdistributed.modelparallel.torch as smp1170 1171 @smp.step()1172 def smp_forward_backward(model, inputs, gradient_accumulation_steps=1):1173 outputs = model(**inputs)1174 loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]1175 loss /= gradient_accumulation_steps1176 model.backward(loss)1177 return loss1178 1179 @smp.step()1180 def smp_forward_only(model, inputs):1181 return model(**inputs)1182 1183 def smp_gather(tensor):1184 if isinstance(tensor, (list, tuple)):1185 return type(tensor)(smp_gather(t) for t in tensor)1186 elif isinstance(tensor, dict):1187 return type(tensor)({k: smp_gather(v) for k, v in tensor.items()})1188 elif not isinstance(tensor, torch.Tensor):1189 raise TypeError(1190 f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors."1191 )1192 all_tensors = smp.allgather(tensor, smp.CommGroup.DP_GROUP)1193 all_tensors = [atleast_1d(t) for t in all_tensors]1194 return torch.cat([t.cpu() for t in all_tensors], dim=0)1195 1196 def smp_nested_concat(tensor):1197 if isinstance(tensor, (list, tuple)):1198 return type(tensor)(smp_nested_concat(t) for t in tensor)1199 elif isinstance(tensor, dict):1200 return type(tensor)({k: smp_nested_concat(v) for k, v in tensor.items()})