softwareweaver/MusicGen

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7from concurrent.futures import ThreadPoolExecutor8from collections import deque9from functools import partial10from hashlib import sha111import logging12from pathlib import Path13import sys14import typing as tp15import zipfile16 17import flashy18import torch19 20 21logger = logging.getLogger(__name__)22 23 24def get_full_embed(full_embed: torch.Tensor, x: tp.Any, idx: int, device: tp.Union[str, torch.device]) -> torch.Tensor:25    """Utility function for the EmbeddingCache, returning the full embedding without any chunking.26    This method can be used in case there is no need in extracting a chunk of the full embedding27    read from the cache.28 29    Args:30        full_embed (torch.Tensor): The full embedding.31        x (any): Batch object from which the full embedding is derived.32        idx (torch.Tensor): Index of object to consider in the batch object.33    Returns:34        full_embed (torch.Tensor): The full embedding35    """36    return full_embed.to(device)37 38 39class EmbeddingCache:40    """Cache around embeddings computation for faster execution.41    The EmbeddingCache is storing pre-computed embeddings on disk and provides a simple API42    to retrieve the pre-computed embeddings on full inputs and extract only a given chunk43    using a user-provided function. When the cache is warm (all embeddings are pre-computed),44    the EmbeddingCache allows for faster training as it removes the need of computing the embeddings.45    Additionally, it provides in-memory cache around the loaded embeddings to limit IO footprint46    and synchronization points in the forward calls.47 48    Args:49        cache_path (Path): Path to folder where all pre-computed embeddings are saved on disk.50        device (str or torch.device): Device on which the embedding is returned.51        compute_embed_fn (callable[[Path, any, int], torch.Tensor], optional): Function to compute52            the embedding from a given object and path. This user provided function can compute the53            embedding from the provided object or using the provided path as entry point. The last parameter54            specify the index corresponding to the current embedding in the object that can represent batch metadata.55        extract_embed_fn (callable[[torch.Tensor, any, int], torch.Tensor], optional): Function to extract56            the desired embedding chunk from the full embedding loaded from the cache. The last parameter57            specify the index corresponding to the current embedding in the object that can represent batch metadata.58            If not specified, will return the full embedding unmodified.59    """60    def __init__(self, cache_path: tp.Union[str, Path], device: tp.Union[str, torch.device],61                 compute_embed_fn: tp.Callable[[Path, tp.Any, int], torch.Tensor],62                 extract_embed_fn: tp.Optional[tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor]] = None):63        self.cache_path = Path(cache_path)64        self.device = device65        self._compute_embed_fn = compute_embed_fn66        self._extract_embed_fn: tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor]67        if extract_embed_fn is not None:68            self._extract_embed_fn = extract_embed_fn69        else:70            self._extract_embed_fn = partial(get_full_embed, device=device)71        if self.cache_path is not None:72            self.cache_path.mkdir(exist_ok=True, parents=True)73            logger.info(f"Cache instantiated at: {self.cache_path}")74            self.pool = ThreadPoolExecutor(8)75            self.pool.__enter__()76        self._current_batch_cache: dict = {}77        self._memory_cache: dict = {}78 79    def _get_cache_path(self, path: tp.Union[Path, str]):80        """Get cache path for the given file path."""81        sig = sha1(str(path).encode()).hexdigest()82        return self.cache_path / sig83 84    @staticmethod85    def _get_full_embed_from_cache(cache: Path):86        """Loads full pre-computed embedding from the cache."""87        try:88            embed = torch.load(cache, 'cpu')89        except Exception as exc:90            logger.error("Error loading %s: %r", cache, exc)91            embed = None92        return embed93 94    def get_embed_from_cache(self, paths: tp.List[Path], x: tp.Any) -> torch.Tensor:95        """Get embedding from cache, computing and storing it to cache if not already cached.96        The EmbeddingCache first tries to load the embedding from the in-memory cache97        containing the pre-computed chunks populated through `populate_embed_cache`.98        If not found, the full embedding is computed and stored on disk to be later accessed99        to populate the in-memory cache, and the desired embedding chunk is extracted and returned.100 101        Args:102            paths (list[Path or str]): List of paths from where the embeddings can be loaded.103            x (any): Object from which the embedding is extracted.104        """105        embeds = []106        for idx, path in enumerate(paths):107            cache = self._get_cache_path(path)108            if cache in self._current_batch_cache:109                embed = self._current_batch_cache[cache]110            else:111                full_embed = self._compute_embed_fn(path, x, idx)112                try:113                    with flashy.utils.write_and_rename(cache, pid=True) as f:114                        torch.save(full_embed.cpu(), f)115                except Exception as exc:116                    logger.error('Error saving embed %s (%s): %r', cache, full_embed.shape, exc)117                else:118                    logger.info('New embed cache saved: %s (%s)', cache, full_embed.shape)119                    embed = self._extract_embed_fn(full_embed, x, idx)120            embeds.append(embed)121        embed = torch.stack(embeds, dim=0)122        return embed123 124    def populate_embed_cache(self, paths: tp.List[Path], x: tp.Any) -> None:125        """Populate in-memory caches for embeddings reading from the embeddings stored on disk.126        The in-memory caches consist in a cache for the full embedding and another cache for the127        final embedding chunk. Such caches are used to limit the IO access when computing the actual embeddings128        and reduce the IO footprint and synchronization points during forward passes.129 130        Args:131            paths (list[Path]): List of paths from where the embeddings can be loaded.132            x (any): Object from which the embedding is extracted.133        """134        self._current_batch_cache.clear()135        if self.cache_path is not None:136            futures: list = []137            for path in paths:138                assert path is not None, "Path is required for computation from cache"139                cache = self._get_cache_path(path)140                if cache in self._memory_cache or not cache.exists():141                    futures.append(None)142                else:143                    futures.append(self.pool.submit(EmbeddingCache._get_full_embed_from_cache, cache))144            for idx, (path, future) in enumerate(zip(paths, futures)):145                assert path is not None146                cache = self._get_cache_path(path)147                full_embed = None148                if future is None:149                    if cache in self._memory_cache:150                        full_embed = self._memory_cache[cache]151                else:152                    full_embed = future.result()153                    if full_embed is not None:154                        self._memory_cache[cache] = full_embed155                        full_embed = full_embed.to(self.device)156                if full_embed is not None:157                    embed = self._extract_embed_fn(full_embed, x, idx)158                    self._current_batch_cache[cache] = embed159 160 161class CachedBatchWriter:162    """Write pre computed caches for mini batches. This can163    make loading a lot more efficient depending on your filesystem.164 165    Args:166        cache_folder (Path): folder in which the cached minibatches167            will be stored.168 169    Inside cache folder, the structure is the following:170    `epoch_number / update_number.zip`171    And the zip file contains one entry per batch item.172 173    It is possible to use the cache with a batch size smaller than174    created with but obviously not larger. Make sure to call the175    `start_epoch(epoch)` method for indicating changes of epochs.176 177    See the grid `audiocraft/grids/musicgen/musicgen_warmup_cache.py`178    for an example of how to warmup the cache.179    """180    def __init__(self, cache_folder: Path):181        self.cache_folder = cache_folder182        self._current_epoch: tp.Optional[int] = None183        self._current_index = 0184 185    def start_epoch(self, epoch: int):186        """Call at the beginning of each epoch.187        """188        self._current_epoch = epoch189        self._current_index = 0190        self._zip_path.parent.mkdir(exist_ok=True, parents=True)191 192    @staticmethod193    def _get_zip_path(cache_folder: Path, epoch: int, index: int):194        return cache_folder / f"{epoch:05d}" / f"{index:06d}.zip"195 196    @property197    def _zip_path(self):198        assert self._current_epoch is not None199        return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, self._current_index)200 201    def save(self, *content):202        """Save one mini batch. This function is distributed-aware203        and will automatically merge all the items from the different204        workers.205        """206        all_contents = []207        for rank in range(flashy.distrib.world_size()):208            their_content = flashy.distrib.broadcast_object(content, src=rank)209            all_contents.append(their_content)210 211        if flashy.distrib.is_rank_zero():212            idx = 0213            with flashy.utils.write_and_rename(self._zip_path) as tmp:214                with zipfile.ZipFile(tmp, 'w') as zf:215                    for content in all_contents:216                        for vals in zip(*content):217                            with zf.open(f'{idx}', 'w') as f:  # type: ignore218                                torch.save(vals, f)219                            idx += 1220        flashy.distrib.barrier()221        self._current_index += 1222 223 224class CachedBatchLoader:225    """Loader for cached mini-batches dumped with `CachedBatchWriter`.226 227    Args:228        cache_folder (Path): folder in which the cached minibatches are stored.229        batch_size (int): batch size (per GPU) expected.230        num_workers (int): number of workers to use for loading.231        min_length (int): minimum expected length for each epoch. If some232            mini-batches are missing, and error is raised.233 234    This is iterable just like a regular DataLoader.235    """236 237    def __init__(self, cache_folder: Path, batch_size: int,238                 num_workers: int = 10, min_length: int = 1):239        self.cache_folder = cache_folder240        self.batch_size = batch_size241        self.num_workers = num_workers242        self.min_length = min_length243        self._current_epoch: tp.Optional[int] = None244        self.sampler = None  # for compatibility with the regular DataLoader245 246    def __len__(self):247        path = CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch or 0, 0).parent248        return len([p for p in path.iterdir() if p.suffix == ".zip"])249 250    def start_epoch(self, epoch: int):251        """Call at the beginning of each epoch.252        """253        self._current_epoch = epoch254 255    def _zip_path(self, index: int):256        assert self._current_epoch is not None257        return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, index)258 259    def _load_one(self, index: int):260        zip_path = self._zip_path(index)261        if not zip_path.exists():262            if index < self.min_length:263                raise RuntimeError(f"Cache should have at least {self.min_length} batches, but {index} doesn't exist")264 265            return None266        mode = "rb" if sys.version_info >= (3, 9) else "r"267        try:268            with zipfile.ZipFile(zip_path, 'r') as zf:269                rank = flashy.distrib.rank()270                world_size = flashy.distrib.world_size()271                root = zipfile.Path(zf)272                items = list(root.iterdir())273                total_batch_size = self.batch_size * world_size274                if len(items) < total_batch_size:275                    raise RuntimeError(276                        f"The cache can handle a max batch size of {len(items)}, "277                        f"but {total_batch_size} is needed.")278                start = rank * self.batch_size279                items = items[start: start + self.batch_size]280                assert len(items) == self.batch_size281                entries = []282                entries = [torch.load(item.open(mode), 'cpu') for item in items]  # type: ignore283                transposed = zip(*entries)284                out = []285                for part in transposed:286                    assert len(part) > 0287                    if isinstance(part[0], torch.Tensor):288                        out.append(torch.stack(part))289                    else:290                        out.append(part)291                return out292        except Exception:293            logger.error("Error when reading zip path %s", zip_path)294            raise295 296    def __iter__(self):297        """This will yields tuples, exactly as provided to the298        `CachedBatchWriter.save` method.299        """300        pool = ThreadPoolExecutor(self.num_workers)301        next_index = 0302        queue = deque()303 304        def _get_next():305            nonlocal next_index306            r = queue.popleft().result()307            if r is None:308                return None309            else:310                queue.append(pool.submit(self._load_one, next_index))311                next_index += 1312            return r313 314        with pool:315            # fill the buffer of fetching jobs.316            for _ in range(2 * self.num_workers):317                queue.append(pool.submit(self._load_one, next_index))318                next_index += 1319            while True:320                batch = _get_next()321                if batch is None:322                    return323                yield batch324