RASMUS/Finnish-ASR-Canary-v2

1#!/usr/bin/env python2# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16import importlib17import math18import os19import sys20from functools import partial21from numbers import Number22from typing import Literal23 24import click25import lightning.pytorch as pl26import torch27from lhotse import compute_num_samples28from omegaconf import OmegaConf29from torch.utils.data import DataLoader, IterableDataset30 31from nemo.collections.speechlm2 import SALM, SALMWithAsrDecoder32from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, MaskType, NeuralType33from nemo.utils import logging34from nemo.utils.trainer_utils import resolve_trainer_cfg35 36 37class ProfilingBatchGenerator:38    """39    ProfilingBatchGenerator is used to generate artificial mini-batches for model training40    and tracking the progress of batch size optimization.41 42    The high-level usage API is the following::43 44        >>> gen = ProfilingBatchGenerator(schema)45        ... finished = False46        ... while not finished:47        ...     batch = gen(input_seq_len, output_seq_len)48        ...     try:49        ...         training_step(model, batch)50        ...         oom = False51        ...     except torch.cuda.OutOfMemoryError:52        ...         oom = True53        ...     finished = gen.advance(oom)54        ... solution = gen.max_batch_size  # The solution of the search problem.55        ... gen.reset()  # Can re-use for other sequence lengths now.56 57    The search terminates once the difference between max working batch size and min OOM batch size58    divided by the latter is smaller than ``rel_gap_thresh`` that difference amounts to a single element.59    For example, a max working batch size is 96 and min OOM batch size is 100 indicates a gap of 0.04,60    which would terminate the search with threshold of 0.05.61 62    In order to generate mini-batches compatible with a given model, the generator:63 64    * accepts a ``schema`` argument in its constructor, and65 66    * accepts input/output sequence lengths in each call to generate a mini-batch.67 68    ``schema`` has the following structure::69 70 71        >>> {72        ...     "cls": tuple | MyBatchType,73        ...     "inputs": [74        ...         {75        ...             "type": NeuralType(...) | Literal["dummy"],76        ...             "seq_length": Literal["input", "output"],77        ...             "vocab_size": int,  # optional, required only for LabelsType78        ...             "name": str,  # optional, indicates kwarg79        ...         },80        ...         ...,81        ...     ]82        ... }83 84    ``cls`` indicates how we should construct the mini-batch. Typically you can just use ``tuple`` for most85    batch schemas. However, if the model expects a specific, e.g., dataclass, you can tell ``ProfilingBatchGenerator``86    to use it. The mini-batch object will be constructed using the items in ``inputs``.87 88    Each element of ``inputs`` specifies a NeMo NeuralType which needs to have a defined ``elements_type``.89    The supported types are ``AudioSignal``, ``LengthsType`` and ``LabelsType``.90    If "type" is not a NeuralType, we interpret that as a placeholder tensor that's not relevant but expected91    by the model/batch constructor. In addition, ``"seq_length"`` key is used to determine whether we should apply92    input or output sequence length to a given tensor.93 94    Optional keys:95 96    * ``vocab_size`` is required for ``LabelsType`` so that we can generate proper label values.97 98    * ``name`` is required if objects of ``cls`` have to be constructed using keyword arguments.99 100    A simple schema example for a model using audio/lengths tensor pair (unsupervised/self-supervised)::101 102        >>> {103        ...     "cls": tuple,104        ...     "inputs": [105        ...         {"type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input"},106        ...         {"type": NeuralType(("B"), LengthsType()), "seq_length": "input"},107        ...     ]108        ... }109 110    """111 112    def __init__(113        self,114        schema: dict,115        start_batch_size: int = 32,116        rel_gap_thresh: float = 0.05,117        device: str = "cuda",118        float_dtype: torch.dtype = torch.float32,119    ):120        self.schema = schema121        self.start_batch_size = start_batch_size122        self.rel_gap_thresh = rel_gap_thresh123        self.device = device124        self.float_dtype = float_dtype125        self.reset()126 127    def __call__(self, input_seq_length: int, output_seq_length: int):128        B = self._current129        select_seq_length = {"input": input_seq_length, "output": output_seq_length}130        batch = []131        names = []132        for item in self.schema["inputs"]:133            nt = item["type"]134            if isinstance(nt, str) and nt == "constant":135                if isinstance(val := item["value"], str) and val == "batch":136                    tnsr = torch.tensor([B], dtype=torch.long, device=self.device)137                else:138                    tnsr = torch.tensor([val], dtype=torch.long, device=self.device)139            elif not isinstance(nt, NeuralType):  # placeholder140                tnsr = torch.tensor([])141            elif isinstance(nt.elements_type, AudioSignal):142                seq_length = select_seq_length[item["seq_length"]]143                tnsr = torch.randn(B, seq_length, dtype=self.float_dtype, device=self.device)144            elif isinstance(nt.elements_type, LengthsType):145                seq_length = select_seq_length[item["seq_length"]]146                tnsr = torch.ones(B, dtype=torch.long, device=self.device) * seq_length147            elif isinstance(nt.elements_type, MaskType):148                seq_length = select_seq_length[item["seq_length"]]149                tnsr = torch.ones(B, seq_length, device=self.device, dtype=torch.bool)150            elif isinstance(nt.elements_type, LabelsType):151                seq_length = select_seq_length[item["seq_length"]]152                tnsr = torch.randint(0, item["vocab_size"], size=(B, seq_length), device=self.device)153            else:154                raise RuntimeError("Unexpected item in oomptimizer schema: {item}")155            batch.append(tnsr)156            names.append(item.get("name"))157        args = [elem for name, elem in zip(names, batch) if name is None]158        kwargs = {name: elem for name, elem in zip(names, batch) if name is not None}159        if not kwargs and self.schema["cls"] == tuple:160            return tuple(args)161        return self.schema["cls"](*args, **kwargs)162 163    @property164    def max_batch_size(self) -> int | None:165        """166        Return the solution of the batch size search problem.167        It will keep returning None until the search is done.168        """169        if (170            self._max_ok is not None171            and self._min_err is not None172            and (self.current_rel_gap <= self.rel_gap_thresh or self._min_err - self._max_ok <= 1)173        ):174            return self._max_ok175        return None176 177    @property178    def current_rel_gap(self) -> float | None:179        """180        Return the current gap between the largest batch that works and the smallest batch that triggers OOM.181        The gap is defined as the batch size difference divided by the larger element.182        E.g., if the best found batch size is 95 and the smallest that triggers OOM is 100, the gap is 0.05.183        """184        if self._min_err is None or self._max_ok is None:185            return None186        return (self._min_err - self._max_ok) / self._min_err187 188    def reset(self):189        """Reset the generator to prepare it for a new search."""190        self._current = self.start_batch_size191        self._max_ok = None  # max batch size that works192        self._min_err = None  # min batch size that doesn't work193 194    def advance(self, oom: bool) -> bool:195        """196        Adjusts the current batch size based on the outcome.197        Returns a bool indicating whether the calibration is complete.198        """199        if self.max_batch_size is not None:200            return True201 202        if oom:203            # Training step failed with OOM.204            # Update the minimum known batch size that causes an error.205            self._min_err = min(float("inf") if self._min_err is None else self._min_err, self._current)206            # Training step failed on OOM207            if self._max_ok is None:208                # We haven't found a batch size that works yet, keep going 2x down.209                self._current = round(self._current / 2)210            else:211                # Try the middle-point between the known extremes.212                self._current = round((self._max_ok + self._min_err) / 2)213        else:214            # Training step successful.215            # Update the maximum known batch size that works.216            self._max_ok = max(-1 if self._max_ok is None else self._max_ok, self._current)217            if self._min_err is None:218                # We haven't found a batch size that causes an error yet, keep going 2x higher219                self._current *= 2220            else:221                # Try the middle-point between the known extremes.222                self._current = round((self._max_ok + self._min_err) / 2)223 224        return False225 226 227class FloatList(click.Option):228    """Support passing bucket duration bins as [1.1,2.5,5.6,...]"""229 230    name = "list[float]"231 232    def type_cast_value(self, ctx, value):233        if isinstance(value, list) and all(isinstance(v, float) for v in value):234            return value235        try:236            import ast237 238            ans = ast.literal_eval(value)239            if isinstance(ans[0], list):240                ans = [tuple(item) for item in ans]241            return ans242        except ValueError:243            raise click.BadParameter(value)244 245 246@click.command(context_settings={'show_default': True})247@click.option(248    "-n",249    "--pretrained-name",250    type=str,251    default=None,252    help="Name of a pretrained model to use, e.g. 'nvidia/canary-1b'.",253)254@click.option(255    "-m",256    "--module-name",257    type=str,258    default=None,259    help="Full path to NeMo's module corresponding to CONFIG_PATH, e.g. 'nemo.collections.asr.models.EncDecMultiTaskModel'.",260)261@click.option(262    "-c", "--config-path", type=str, default=None, help="Path to the training configuration file for MODULE_NAME."263)264@click.option(265    "-b",266    "--buckets",267    cls=FloatList,268    default=[5.0, 10.0, 15.0, 20.0, 25.0, 30.0],269    help="List of upper-bound bucket bins (i.e. first bucket is [0.0 - item0), second bucket is [item0 - item1), etc.). "270    "We also support a nested list for 2D bucketing, e.g. [[2.0, 10],[2.0,20],[4.5,15],[4.5,30],...], "271    "where each item is a pair of (max_input_seq_len, max_output_seq_len) for a given bucket.",272)273@click.option(274    "-t",275    "--threshold",276    type=float,277    default=0.05,278    help="Search stopping criterion in range [0, 1], lower is more precise. Interpret as the uncerainty gap, i.e. (min_oom_batch_size - max_ok_batch_size) / min_oom_batch_size.",279)280@click.option("-s", "--start-batch-size", type=int, default=32, help="Initial batch size to start the search from.")281@click.option(282    "-r",283    "--ratio",284    type=float,285    default=12,  # conservative estimate towards longer transcripts286    help="The output_sequence_length to input_sequence_length ratio for the purpose of determing the maximum output sequence lengths. "287    "The interpretation depends on input and output modalities. Examples: for audio->text it's tokens per second. "288    "For text->audio it's seconds per token. For audio->audio it's output seconds per input second. "289    "For text->text it's output tokens per input token. "290    "In general larger ratio means longer output sequences and increased memory consumption. "291    "The default value is set adequately for automatic speech recognition. "292    "This argument is ignored when 2D buckets are provided to --buckets option.",293)294@click.option(295    "-f",296    "--memory-fraction",297    type=float,298    default=0.9,299    help="Limits the use of CUDA memory for this process to MEMORY_FRACTION of the total device memory. "300    "By default we force 5% memory to be unused to account for non-training-loop related CUDA memory usage"301    "in actual training scripts.",302)303@click.option(304    "-y",305    "--dtype",306    default="bfloat16",307    help="Float precision to use for computation (used together with autocast).",308)309@click.option(310    "--ddp/--no-ddp",311    type=bool,312    default=True,313    help="Whether we should simulate DDP GPU RAM usage. Stores an extra copy of the model in GPU memory. Enabled by default.",314)315def oomptimizer(316    pretrained_name: str | None,317    module_name: str | None,318    config_path: str | None,319    buckets: list[float],320    threshold: float,321    start_batch_size: int,322    ratio: float,323    memory_fraction: float,324    dtype: str,325    ddp: bool,326):327    """328    OOMptimizer finds the optimal batch sizes for training your model with bucketing dataloading.329    It performs a search over batch sizes until it converges by measuring the GPU memory usage for330    a model's training step and optimizer update.331 332    \b333    There are two main usage patterns: for using a pretrained model or an untrained model configuration.334    The latter is more flexible but requires the user to provide two separate arguments. Examples:335    * python oomptimizer.py --pretrained-name nvidia/canary-1b336    * python oomptimizer.py --module-name nemo.collections.asr.models.EncDecMultiTaskModel \337        --config-path examples/asr/conf/speech_multitask/fast-conformer_aed.yaml338 339    Dynamic bucketing is notoriously difficult to tune as you risk running into CUDA OOM many steps into the training.340    In order to simplify finding the optimal settings, OOMptimizer scans each bucket to find the maximum possible341    batch size that doesn't trigger a CUDA OOM.342 343    \b344    The suggested workflow is the following:345    1) Run scripts/speech_recognition/estimate_duration_bins.py to get the duration distribution of your data.346        (consider running estimate_duration_bins_2d.py for models with a strong dependency on output sequence length347        such as attention-encoder-decoder models).348    2) Run OOMptimizer to find the optimal batch sizes for your specific model, optimizer, and GPU.349    3) Use these optimal settings in your actual training script and enjoy optimal GPU utilization OOM-free.350 351    In the unlikely event that OOMptimizer bucket batch sizes are still leading to OOMs,352    please try a lower setting of the MEMORY_FRACTION option, e.g. 0.75 (75% of GPU memory).353    This may be required in very complex setups where there are additional GPU RAM loads that can't be anticipated354    through the combination of training_step and optimizer update.355    """356    assert pretrained_name is None, "--pretrained-name is not supported yet for Duplex S2S"357    if all(opt is None for opt in (pretrained_name, module_name, config_path)):358        click.secho(359            "You need to provide either PRETRAINED_NAME or the pair of MODULE_NAME and CONFIG_PATH.", fg="yellow"360        )361        sys.exit(1)362    logging.setLevel(logging.CRITICAL)363    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))364    device = torch.device(f'cuda:{os.environ["LOCAL_RANK"]}')365    dtype = getattr(torch, dtype)366    torch.cuda.set_per_process_memory_fraction(memory_fraction, device)367 368    torch.distributed.init_process_group(backend="nccl")369    torch.set_float32_matmul_precision("medium")370    torch.backends.cudnn.allow_tf32 = True371 372    assert config_path is not None, "--module-name requires --config-path to be specified as well."373    assert module_name is not None, "--config-path requires --module-name to be specified as well."374    cfg = OmegaConf.load(config_path)375    namespace, name = module_name.rsplit('.', maxsplit=1)376    model_cls = getattr(importlib.import_module(namespace), name)377    trainer = pl.Trainer(378        **{379            **resolve_trainer_cfg(cfg.trainer),380            "max_steps": 1,381            "max_epochs": 1,382            "limit_val_batches": 0.0,383            "val_check_interval": 0.0,384        }385    )386    with trainer.init_module():387        model = model_cls(OmegaConf.to_container(cfg.model, resolve=True))388    model = model.to(device)389 390    if isinstance(model, (SALM, SALMWithAsrDecoder)):391        model.prepare_inputs = partial(_override_prepare_inputs, model)392 393    if not hasattr(model, "oomptimizer_schema"):394        click.secho(395            f"We read model of type {type(model)} which doesn't seem to support OOMptimizer "396            f"(we could not find the property .oomptimizer_schema).",397            fg="red",398        )399        sys.exit(1)400 401    schema = model.oomptimizer_schema402 403    is_2d_bucketing = all(404        isinstance(item, (list, tuple)) and len(item) == 2 and all(isinstance(v, Number) for v in item)405        for item in buckets406    )407    # Determine modality for input and output.408    modalities = [409        (410            "text"411            if any(412                isinstance(item["type"], NeuralType)413                and isinstance(item["type"].elements_type, LabelsType)414                and item["seq_length"] == direction415                for item in schema["inputs"]416                if item["type"] != "dummy"417            )418            else "audio"419        )420        for direction in ("input", "output")421    ]422 423    def get_max_seq_lens(buckets):424 425        def _determine_lens_for_bucket(bin):426            if isinstance(model, (SALM, SALMWithAsrDecoder)):427                return bin, bin  # Note: only 1D bucketing, only counted in tokens428            elif is_2d_bucketing:429                input_len, output_len = bin430            else:431                input_len = bin432                output_len = math.ceil(ratio * input_len)433            sampling_rate = getattr(434                model, "sample_rate", 16000435            )  # TODO: may need to extend schema for broader model coverage436            match modalities:437                case "audio", "audio":438                    return (439                        compute_num_samples(input_len, sampling_rate=sampling_rate),440                        compute_num_samples(output_len, sampling_rate=sampling_rate),441                    )442                case "audio", "text":443                    return (compute_num_samples(input_len, sampling_rate=sampling_rate), output_len)444                case "text", "audio":445                    return (446                        input_len,447                        compute_num_samples(output_len, sampling_rate=sampling_rate),448                    )449                case "text", "text":450                    return input_len, output_len451                case _:452                    raise RuntimeError(f"Unexpected modality combination: {_}")453 454        return [_determine_lens_for_bucket(bin) for bin in buckets]455 456    click.echo("Starting profiling.")457    max_seq_lens = get_max_seq_lens(buckets)458    gen = ProfilingBatchGenerator(459        schema=schema, start_batch_size=start_batch_size, rel_gap_thresh=threshold, device=device, float_dtype=dtype460    )461    profile = {}462 463    class _GenDataset(IterableDataset):464        def __iter__(self):465            gen.reset()466            gen._current = 1467            yield gen(33, 33)468            # yield gen(16000, 13)469            gen.reset()470 471        def __len__(self):472            return 1473 474    # initialize everything PTL needs475    trainer.fit(model, DataLoader(_GenDataset(), batch_size=None))476    model = model.to(device)477    optimizer = model.configure_optimizers()["optimizer"]478    model.log = lambda *args, **kwargs: None  # no logging479 480    # Iterate buckets from the largest to the smallest sequences. This usually ends up creating481    # a tiny bit smaller batches, likely due to worse memory fragmentation.482    with torch.autocast("cuda", dtype=None, enabled=False):483        for bucket, (seq_len_in, seq_len_out) in reversed(list(zip(buckets, max_seq_lens))):484            click.echo(f"The current sequence lengths are: input={seq_len_in} output={seq_len_out}.")485            gen.reset()486            batch_idx = 0487 488            def step():489                click.echo(490                    f"\t[BEGIN step] [CUDA RAM CURRENT: {torch.cuda.memory_allocated() / (1024 * 1024):.1f}MB] [CUDA RAM MAX: {torch.cuda.max_memory_allocated() / (1024*1024):.1f}MB]"491                )492                batch = gen(seq_len_in, seq_len_out)493 494                oom = False495                try:496                    click.echo(f"\tCurrent gap: {gen.current_rel_gap}... ", nl=False)497                    optimizer.zero_grad()498                    out = model.training_step(batch, batch_idx)499                    out['loss'].sum().backward()500                    optimizer.step()501                except torch.cuda.OutOfMemoryError as e:502                    click.secho(f"OOM!", fg="yellow")503                    oom = True504                except RuntimeError as e:505                    if "cuFFT error: CUFFT_INTERNAL_ERROR" not in str(e):506                        raise507                    click.secho(f"OOM!", fg="yellow")508                    oom = True509                else:510                    click.secho(f"OK!", fg="green")511                finally:512                    click.echo(513                        f"\t[END step] [CUDA RAM CURRENT: {torch.cuda.memory_allocated() / (1024 * 1024):.1f}MB] [CUDA RAM MAX: {torch.cuda.max_memory_allocated() / (1024*1024):.1f}MB]"514                    )515                    del batch516                    # Note: We could call empty_cache() to free up some more memory on the GPU,517                    #       but we have found out empirically that this causes a mismatched condition518                    #       between OOMptimizer and the actual training. During training, there is some519                    #       degree of memory fragmentation and it's better to simulate that in OOMptimizer.520                    # torch.cuda.memory.empty_cache()521                    torch.cuda.reset_max_memory_allocated()522                return oom523 524            oom = step()525            while not (finished := gen.advance(oom)):526                click.echo("\t" + "=" * 80)527                oom = step()528 529            click.secho(530                f"=> Optimal setting for bucket={bucket} (input={seq_len_in} output={seq_len_out}) is max_batch_size={gen.max_batch_size}",531                fg="green",532            )533            profile[(bucket, seq_len_in, seq_len_out)] = gen.max_batch_size534            gen.start_batch_size = gen.max_batch_size * 2535 536    # Reverse the profile to be ascendingly sorted again.537    profile = dict(reversed(list(profile.items())))538 539    click.echo("The 1st stage profile is:")540    for (bucket, seq_len_in, seq_len_out), bs in profile.items():541        click.echo(f"Bucket={bucket} (input={seq_len_in} output={seq_len_out}) => max_batch_size={bs}")542 543    if is_2d_bucketing:544        # 2D bucketing doesn't support bucket merging.545        final_profile = [["[" + ",".join(map(str, b)) + "]", bs] for (b, _, __), bs in profile.items()]546    else:547        click.echo("Bucket merging stage...")548        final_profile = []549        for idx, ((bucket, seq_len_in, seq_len_out), bs) in enumerate(profile.items()):550            if idx == 0:551                final_profile.append([bucket, bs])552                continue553            if bs == final_profile[-1][1]:554                click.echo(f"Merging bucket {idx} with bucket {idx-1} due to identical batch sizes.")555                final_profile[-1][0] = bucket556                continue557            final_profile.append([bucket, bs])558 559    click.secho(f"The profile was created with the following settings:")560    click.secho(f"* using {memory_fraction:.1%} of available GPU RAM.")561    click.secho(f"* {'' if ddp else 'not '}simulating DDP memory overhead.")562    click.secho(f"* using AMP with dtype={dtype}.")563    click.secho("The final profile is:", bold=True)564    click.secho("\tbucket_duration_bins=[" + ",".join(str(seqlen) for seqlen, bs in final_profile) + "]", bold=True)565    click.secho("\tbucket_batch_size=[" + ",".join(str(bs) for seqlen, bs in final_profile) + "]", bold=True)566 567 568def _override_prepare_inputs(self, batch: dict) -> dict:569    ratio = 0.8570    input_embs = self.embed_tokens(batch["input_ids"][:, :-1])571    target_ids = batch["input_ids"][:, 1:]572    attention_mask = torch.ones_like(target_ids, dtype=torch.bool)573 574    B, T = input_embs.shape[:2]575    audio_emb_len = int(input_embs.shape[1] * ratio)576    n_samples = int(audio_emb_len * self.token_equivalent_duration * self.sampling_rate)577    audio = torch.randn(B, n_samples, device=input_embs.device, dtype=torch.float32)578    audio_lens = torch.tensor([n_samples] * B, device=input_embs.device)579    audio_embs, _ = self.perception(input_signal=audio, input_signal_length=audio_lens)580    input_embs[:, : audio_embs.shape[1]] = audio_embs581 582    return {583        "input_embeds": input_embs,584        "attention_mask": attention_mask,585        "target_ids": target_ids,586    }587 588 589if __name__ == "__main__":590    oomptimizer()591