Spaces:
Build error
Build error
File size: 14,289 Bytes
00f2826 |
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 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from concurrent.futures import ThreadPoolExecutor
from collections import deque
from functools import partial
from hashlib import sha1
import logging
from pathlib import Path
import sys
import typing as tp
import zipfile
import flashy
import torch
logger = logging.getLogger(__name__)
def get_full_embed(full_embed: torch.Tensor, x: tp.Any, idx: int, device: tp.Union[str, torch.device]) -> torch.Tensor:
"""Utility function for the EmbeddingCache, returning the full embedding without any chunking.
This method can be used in case there is no need in extracting a chunk of the full embedding
read from the cache.
Args:
full_embed (torch.Tensor): The full embedding.
x (any): Batch object from which the full embedding is derived.
idx (torch.Tensor): Index of object to consider in the batch object.
Returns:
full_embed (torch.Tensor): The full embedding
"""
return full_embed.to(device)
class EmbeddingCache:
"""Cache around embeddings computation for faster execution.
The EmbeddingCache is storing pre-computed embeddings on disk and provides a simple API
to retrieve the pre-computed embeddings on full inputs and extract only a given chunk
using a user-provided function. When the cache is warm (all embeddings are pre-computed),
the EmbeddingCache allows for faster training as it removes the need of computing the embeddings.
Additionally, it provides in-memory cache around the loaded embeddings to limit IO footprint
and synchronization points in the forward calls.
Args:
cache_path (Path): Path to folder where all pre-computed embeddings are saved on disk.
device (str or torch.device): Device on which the embedding is returned.
compute_embed_fn (callable[[Path, any, int], torch.Tensor], optional): Function to compute
the embedding from a given object and path. This user provided function can compute the
embedding from the provided object or using the provided path as entry point. The last parameter
specify the index corresponding to the current embedding in the object that can represent batch metadata.
extract_embed_fn (callable[[torch.Tensor, any, int], torch.Tensor], optional): Function to extract
the desired embedding chunk from the full embedding loaded from the cache. The last parameter
specify the index corresponding to the current embedding in the object that can represent batch metadata.
If not specified, will return the full embedding unmodified.
"""
def __init__(self, cache_path: tp.Union[Path], device: tp.Union[str, torch.device],
compute_embed_fn: tp.Callable[[Path, tp.Any, int], torch.Tensor],
extract_embed_fn: tp.Optional[tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor]] = None):
self.cache_path = Path(cache_path)
self.device = device
self._compute_embed_fn = compute_embed_fn
self._extract_embed_fn: tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor]
if extract_embed_fn is not None:
self._extract_embed_fn = extract_embed_fn
else:
self._extract_embed_fn = partial(get_full_embed, device=device)
if self.cache_path is not None:
self.cache_path.mkdir(exist_ok=True, parents=True)
logger.info(f"Cache instantiated at: {self.cache_path}")
self.pool = ThreadPoolExecutor(8)
self.pool.__enter__()
self._current_batch_cache: dict = {}
self._memory_cache: dict = {}
def _get_cache_path(self, path: tp.Union[Path, str]):
"""Get cache path for the given file path."""
sig = sha1(str(path).encode()).hexdigest()
return self.cache_path / sig
@staticmethod
def _get_full_embed_from_cache(cache: Path):
"""Loads full pre-computed embedding from the cache."""
try:
embed = torch.load(cache, 'cpu')
except Exception as exc:
logger.error("Error loading %s: %r", cache, exc)
embed = None
return embed
def get_embed_from_cache(self, paths: tp.List[Path], x: tp.Any) -> torch.Tensor:
"""Get embedding from cache, computing and storing it to cache if not already cached.
The EmbeddingCache first tries to load the embedding from the in-memory cache
containing the pre-computed chunks populated through `populate_embed_cache`.
If not found, the full embedding is computed and stored on disk to be later accessed
to populate the in-memory cache, and the desired embedding chunk is extracted and returned.
Args:
paths (list[Path or str]): List of paths from where the embeddings can be loaded.
x (any): Object from which the embedding is extracted.
"""
embeds = []
for idx, path in enumerate(paths):
cache = self._get_cache_path(path)
if cache in self._current_batch_cache:
embed = self._current_batch_cache[cache]
else:
full_embed = self._compute_embed_fn(path, x, idx)
try:
with flashy.utils.write_and_rename(cache, pid=True) as f:
torch.save(full_embed.cpu(), f)
except Exception as exc:
logger.error('Error saving embed %s (%s): %r', cache, full_embed.shape, exc)
else:
logger.info('New embed cache saved: %s (%s)', cache, full_embed.shape)
embed = self._extract_embed_fn(full_embed, x, idx)
embeds.append(embed)
embed = torch.stack(embeds, dim=0)
return embed
def populate_embed_cache(self, paths: tp.List[Path], x: tp.Any) -> None:
"""Populate in-memory caches for embeddings reading from the embeddings stored on disk.
The in-memory caches consist in a cache for the full embedding and another cache for the
final embedding chunk. Such caches are used to limit the IO access when computing the actual embeddings
and reduce the IO footprint and synchronization points during forward passes.
Args:
paths (list[Path]): List of paths from where the embeddings can be loaded.
x (any): Object from which the embedding is extracted.
"""
self._current_batch_cache.clear()
if self.cache_path is not None:
futures: list = []
for path in paths:
assert path is not None, "Path is required for computation from cache"
cache = self._get_cache_path(path)
if cache in self._memory_cache or not cache.exists():
futures.append(None)
else:
futures.append(self.pool.submit(EmbeddingCache._get_full_embed_from_cache, cache))
for idx, (path, future) in enumerate(zip(paths, futures)):
assert path is not None
cache = self._get_cache_path(path)
full_embed = None
if future is None:
if cache in self._memory_cache:
full_embed = self._memory_cache[cache]
else:
full_embed = future.result()
if full_embed is not None:
self._memory_cache[cache] = full_embed
full_embed = full_embed.to(self.device)
if full_embed is not None:
embed = self._extract_embed_fn(full_embed, x, idx)
self._current_batch_cache[cache] = embed
class CachedBatchWriter:
"""Write pre computed caches for mini batches. This can
make loading a lot more efficient depending on your filesystem.
Args:
cache_folder (Path): folder in which the cached minibatches
will be stored.
Inside cache folder, the structure is the following:
`epoch_number / update_number.zip`
And the zip file contains one entry per batch item.
It is possible to use the cache with a batch size smaller than
created with but obviously not larger. Make sure to call the
`start_epoch(epoch)` method for indicating changes of epochs.
See the grid `audiocraft/grids/musicgen/musicgen_warmup_cache.py`
for an example of how to warmup the cache.
"""
def __init__(self, cache_folder: Path):
self.cache_folder = cache_folder
self._current_epoch: tp.Optional[int] = None
self._current_index = 0
def start_epoch(self, epoch: int):
"""Call at the beginning of each epoch.
"""
self._current_epoch = epoch
self._current_index = 0
self._zip_path.parent.mkdir(exist_ok=True, parents=True)
@staticmethod
def _get_zip_path(cache_folder: Path, epoch: int, index: int):
return cache_folder / f"{epoch:05d}" / f"{index:06d}.zip"
@property
def _zip_path(self):
assert self._current_epoch is not None
return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, self._current_index)
def save(self, *content):
"""Save one mini batch. This function is distributed-aware
and will automatically merge all the items from the different
workers.
"""
all_contents = []
for rank in range(flashy.distrib.world_size()):
their_content = flashy.distrib.broadcast_object(content, src=rank)
all_contents.append(their_content)
if flashy.distrib.is_rank_zero():
idx = 0
with flashy.utils.write_and_rename(self._zip_path) as tmp:
with zipfile.ZipFile(tmp, 'w') as zf:
for content in all_contents:
for vals in zip(*content):
with zf.open(f'{idx}', 'w') as f: # type: ignore
torch.save(vals, f)
idx += 1
flashy.distrib.barrier()
self._current_index += 1
class CachedBatchLoader:
"""Loader for cached mini-batches dumped with `CachedBatchWriter`.
Args:
cache_folder (Path): folder in which the cached minibatches are stored.
batch_size (int): batch size (per GPU) expected.
num_workers (int): number of workers to use for loading.
min_length (int): minimum expected length for each epoch. If some
mini-batches are missing, and error is raised.
This is iterable just like a regular DataLoader.
"""
def __init__(self, cache_folder: Path, batch_size: int,
num_workers: int = 10, min_length: int = 1):
self.cache_folder = cache_folder
self.batch_size = batch_size
self.num_workers = num_workers
self.min_length = min_length
self._current_epoch: tp.Optional[int] = None
self.sampler = None # for compatibility with the regular DataLoader
def __len__(self):
path = CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch or 0, 0).parent
return len([p for p in path.iterdir() if p.suffix == ".zip"])
def start_epoch(self, epoch: int):
"""Call at the beginning of each epoch.
"""
self._current_epoch = epoch
def _zip_path(self, index: int):
assert self._current_epoch is not None
return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, index)
def _load_one(self, index: int):
zip_path = self._zip_path(index)
if not zip_path.exists():
if index < self.min_length:
raise RuntimeError(f"Cache should have at least {self.min_length} batches, but {index} doesn't exist")
return None
mode = "rb" if sys.version_info >= (3, 9) else "r"
try:
with zipfile.ZipFile(zip_path, 'r') as zf:
rank = flashy.distrib.rank()
world_size = flashy.distrib.world_size()
root = zipfile.Path(zf)
items = list(root.iterdir())
total_batch_size = self.batch_size * world_size
if len(items) < total_batch_size:
raise RuntimeError(
f"The cache can handle a max batch size of {len(items)}, "
f"but {total_batch_size} is needed.")
start = rank * self.batch_size
items = items[start: start + self.batch_size]
assert len(items) == self.batch_size
entries = []
entries = [torch.load(item.open(mode), 'cpu') for item in items] # type: ignore
transposed = zip(*entries)
out = []
for part in transposed:
assert len(part) > 0
if isinstance(part[0], torch.Tensor):
out.append(torch.stack(part))
else:
out.append(part)
return out
except Exception:
logger.error("Error when reading zip path %s", zip_path)
raise
def __iter__(self):
"""This will yields tuples, exactly as provided to the
`CachedBatchWriter.save` method.
"""
pool = ThreadPoolExecutor(self.num_workers)
next_index = 0
queue = deque()
def _get_next():
nonlocal next_index
r = queue.popleft().result()
if r is None:
return None
else:
queue.append(pool.submit(self._load_one, next_index))
next_index += 1
return r
with pool:
# fill the buffer of fetching jobs.
for _ in range(2 * self.num_workers):
queue.append(pool.submit(self._load_one, next_index))
next_index += 1
while True:
batch = _get_next()
if batch is None:
return
yield batch
|