File size: 8,505 Bytes
d9115cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 | import warnings
from time import sleep
from typing import Any, Dict, List, Optional, Tuple, Union
import torch
from omegaconf import DictConfig
from transformers.cache_utils import DynamicCache
from src.utils import pylogger, rich_utils
from functools import partial
from tokenizers.processors import TemplateProcessing
from src.data.loading.components.interfaces import TokenizerConfig
log = pylogger.RankedLogger(__name__, rank_zero_only=True)
def print_warnings_for_missing_configs(cfg: DictConfig) -> None:
_DEFAULT_CONFIGS = [
"data_loading",
"model",
"loss",
"optim",
"eval",
]
has_warnings = False
for config in _DEFAULT_CONFIGS:
if not cfg.get(config):
log.warning(
f"Config {config} was not found in the config tree. Make sure this is expected."
)
has_warnings = True
if has_warnings:
sleep(3) # wait for 3 seconds to let the user read the warning
def extras(cfg: DictConfig) -> None:
"""Applies optional utilities before the task is started.
Utilities:
- Ignoring python warnings
- Setting tags from command line
- Rich config printing
:param cfg: A DictConfig object containing the config tree.
"""
# return if no `extras` config
if not cfg.get("extras"):
log.warning("Extras config not found! <cfg.extras=null>")
return
# disable python warnings
if cfg.extras.get("ignore_warnings"):
log.info("Disabling python warnings! <cfg.extras.ignore_warnings=True>")
warnings.filterwarnings("ignore")
# prompt user to input tags from command line if none are provided in the config
if cfg.extras.get("enforce_tags"):
log.info("Enforcing tags! <cfg.extras.enforce_tags=True>")
rich_utils.enforce_tags(cfg, save_to_file=True)
if cfg.extras.get("print_config_warnings"):
print_warnings_for_missing_configs(cfg)
# pretty print config tree using Rich library
if cfg.extras.get("print_config"):
log.info("Printing config tree with Rich! <cfg.extras.print_config=True>")
rich_utils.print_config_tree(cfg, resolve=True, save_to_file=True)
def delete_module(module: torch.nn.Module, module_name: str) -> None:
"""Recursively delete a submodule from a module.
:param module: the parent module that we want the submodule to be removed from.
:param module_name: the name of the submodule to be removed.
:return: None.
"""
if hasattr(module, module_name):
delattr(module, module_name)
for name, submodule in module.named_children():
delete_module(submodule, module_name)
def find_module_shape(
module: torch.nn.Module, module_name: str
) -> Optional[torch.Size]:
"""Recursively find a submodule in a module and return its shape.
:param module: the parent module that we want the submodule to be removed from.
:param module_name: the name of the submodule to be removed.
:return: the shape of the module if it exists.
"""
if hasattr(module, module_name):
return getattr(module, module_name).weight.shape
for name, submodule in module.named_children():
shape = find_module_shape(submodule, module_name)
if shape:
return shape
return None
def reset_parameters(module: torch.nn.Module) -> None:
"""Reset the parameters of a given module.
:param module: the module whose parameters will be reset.
:return: None.
"""
if hasattr(module, "reset_parameters"):
module.reset_parameters()
else:
for layer in module.children():
reset_parameters(layer)
def get_var_if_not_none(value: Optional[Any], default_value: Any) -> Any:
"""
:return: value if value is not None, else default_value
Note that when value is:
Boolean: False is not None
Int: 0 is not None
List: An empty list is not None
Tensor: A tensor with all zeros is not None
"""
return value if value is not None else default_value
def get_class_name_str(class_definition: Any) -> str:
"""
Args:
class_definition: The class definition.
Returns:
The fully qualified name of the given class.
"""
return ".".join([class_definition.__module__, class_definition.__name__])
def has_class_object_inside_list(obj_list: list, class_type: Any) -> bool:
"""
Args:
obj_list: List of objects.
class_type: The class type to check.
Returns:
True if the list contains an object of the given class type.
"""
return any(isinstance(obj, class_type) for obj in obj_list)
def convert_legacy_kv_cache_to_dynamic(
kv_cache: Union[DynamicCache, Tuple[torch.Tensor]]
) -> DynamicCache:
"""
Converts a legacy key-value cache (Tuple of tensors) to a dynamic cache.
Args:
kv_cache: The key-value cache.
Returns:
The dynamic cache.
"""
if isinstance(kv_cache, DynamicCache):
return kv_cache
return DynamicCache.from_legacy_cache(kv_cache)
def get_parent_module_and_attr(
model: torch.nn.Module, module_name: str
) -> Tuple[torch.nn.Module, str]:
"""
Get the parent module and attribute name for a given module name.
Args:
model (torch.nn.Module): The model containing the module.
module_name (str): The full name of the module.
Returns:
Tuple[torch.nn.Module, str]: The parent module and the attribute name.
"""
parts = module_name.split(".")
parent = model
for part in parts[:-1]:
parent = getattr(parent, part)
return parent, parts[-1]
def lightning_precision_to_dtype(precision: str) -> torch.dtype:
"""
Convert Lightning precision identifier to PyTorch dtype.
Args:
precision (str): The precision identifier used in Lightning.
Expected values include '32', '32-true', '16', '16-mixed', 'bf16', '64', 'half'.
Returns:
torch.dtype: The corresponding PyTorch dtype.
Raises:
ValueError: If an unsupported precision type is provided.
"""
# Mapping from Lightning precision identifiers to PyTorch dtypes
precision_map = {
"32": torch.float32, # Single precision (float32)
"32-true": torch.float32, # Also maps to float32, useful for clarity when specifying defaults
"64": torch.float64, # Double precision
"16": torch.float16, # Half precision
"16-mixed": torch.float16, # Mixed precision typically uses torch.float16
"bf16": torch.bfloat16, # BFloat16 precision
"half": torch.float16, # Alias for half precision
}
if precision in precision_map:
return precision_map[precision]
else:
raise ValueError(
f"Unsupported precision type: '{precision}'. "
"Supported precision types are: '32', '32-true', '64', '16', '16-mixed', 'bf16', 'half'."
)
def load_tokenize(config: TokenizerConfig) -> Any:
"""Load tokenizer and return a partial function for tokenization."""
tokenizer = config.tokenizer
if hasattr(config, "special_tokens"):
tokenizer.add_special_tokens(config.special_tokens)
if config.postprocess_eos_token:
tokenizer._tokenizer.post_processor = TemplateProcessing(
single="$A " + tokenizer.eos_token,
special_tokens=[(tokenizer.eos_token, tokenizer.eos_token_id)],
)
tokenize = partial(
tokenizer.encode_plus,
max_length=config.max_length,
padding=config.padding,
truncation=config.truncation,
add_special_tokens=config.add_special_tokens,
return_tensors="pt",
)
return tokenize
def sample_gumbel(shape: Tuple, device: torch.device, eps=1e-20) -> torch.Tensor:
"""Sample from Gumbel(0, 1)"""
U = torch.rand(shape, device=device)
return -torch.log(-torch.log(U + eps) + eps)
def gumbel_softmax_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor:
"""Draw a sample from the Gumbel-Softmax distribution"""
y = logits + sample_gumbel(logits.shape, logits.device)
sample = F.softmax(y / temperature, dim=-1)
return sample
|