|
import logging |
|
import math |
|
import os |
|
from contextlib import contextmanager |
|
|
|
import timm.models.hub as timm_hub |
|
import torch |
|
import torch.distributed as dist |
|
import torch.nn as nn |
|
from transformers import StoppingCriteria, StoppingCriteriaList |
|
|
|
|
|
def is_dist_avail_and_initialized(): |
|
if not dist.is_available(): |
|
return False |
|
if not dist.is_initialized(): |
|
return False |
|
return True |
|
|
|
|
|
def get_rank(): |
|
if not is_dist_avail_and_initialized(): |
|
return 0 |
|
return dist.get_rank() |
|
|
|
|
|
def is_main_process(): |
|
return get_rank() == 0 |
|
|
|
|
|
def rank0_print(msg, **kwargs): |
|
if is_main_process(): |
|
print(msg, **kwargs) |
|
|
|
|
|
def download_cached_file(url, check_hash=True, progress=False): |
|
""" |
|
Download a file from a URL and cache it locally. If the file already exists, it is not downloaded again. |
|
If distributed, only the main process downloads the file, and the other processes wait for the file to be downloaded. |
|
""" |
|
def get_cached_file_path(): |
|
|
|
parts = torch.hub.urlparse(url) |
|
filename = os.path.basename(parts.path) |
|
cached_file = os.path.join(timm_hub.get_cache_dir(), filename) |
|
|
|
return cached_file |
|
|
|
if is_main_process(): |
|
timm_hub.download_cached_file(url, check_hash, progress) |
|
|
|
if is_dist_avail_and_initialized(): |
|
dist.barrier() |
|
|
|
return get_cached_file_path() |
|
|
|
|
|
@contextmanager |
|
def all_logging_disabled(highest_level=logging.CRITICAL): |
|
""" |
|
A context manager that will prevent any logging messages |
|
triggered during the body from being processed. |
|
:param highest_level: the maximum logging level in use. |
|
This would only need to be changed if a custom level greater than CRITICAL |
|
is defined. |
|
""" |
|
|
|
|
|
|
|
|
|
|
|
previous_level = logging.root.manager.disable |
|
|
|
logging.disable(highest_level) |
|
|
|
try: |
|
yield |
|
finally: |
|
logging.disable(previous_level) |
|
|
|
|
|
class LoRALinear(nn.Linear): |
|
def __init__(self, |
|
in_features: int, |
|
out_features: int, |
|
bias: bool = True, |
|
device=None, |
|
dtype=None, |
|
lora_r=8, |
|
lora_alpha=16, |
|
lora_dropout=0.05, |
|
**kwargs) -> None: |
|
super().__init__(in_features, out_features, bias, device, dtype) |
|
self.lora_r = lora_r |
|
self.lora_alpha = lora_alpha |
|
if lora_dropout > 0.0: |
|
self.lora_dropout = nn.Dropout(p=lora_dropout) |
|
else: |
|
self.lora_dropout = lambda x: x |
|
self.lora_scaling = self.lora_alpha / self.lora_r |
|
|
|
self.lora_A = nn.Linear(in_features, |
|
self.lora_r, |
|
bias=False, |
|
device=device, |
|
dtype=dtype) |
|
self.lora_B = nn.Linear(self.lora_r, |
|
out_features, |
|
bias=False, |
|
device=device, |
|
dtype=dtype) |
|
|
|
self.reset_parameters() |
|
|
|
def reset_parameters(self): |
|
if hasattr(self, "lora_A"): |
|
|
|
nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) |
|
nn.init.zeros_(self.lora_B.weight) |
|
|
|
def forward(self, x): |
|
orig_type = x.dtype |
|
res = super().forward(x) |
|
x = x.float() |
|
res += self.lora_B(self.lora_A( |
|
self.lora_dropout(x))) * self.lora_scaling |
|
return res.to(orig_type) |
|
|
|
|
|
class StoppingCriteriaSub(StoppingCriteria): |
|
def __init__(self, stops=[], encounters=1): |
|
super().__init__() |
|
self.stops = stops |
|
|
|
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor): |
|
for stop in self.stops: |
|
if torch.all((stop == input_ids[:, -len(stop):])).item(): |
|
return True |
|
|
|
return False |
|
|
