RASMUS/Finnish-ASR-Canary-v2

1#!/usr/bin/env python2# Copyright (c) 2025, 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 sys19from numbers import Number20from typing import Iterable, Literal21 22import click23import lightning.pytorch as pl24import torch25from lhotse import compute_num_samples26from omegaconf import OmegaConf27 28from nemo.collections.asr.models.asr_model import ASRModel29from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, MaskType, NeuralType30from nemo.utils import logging31 32 33class ProfilingBatchGenerator:34    """35    ProfilingBatchGenerator is used to generate artificial mini-batches for model training36    and tracking the progress of batch size optimization.37 38    The high-level usage API is the following::39 40        >>> gen = ProfilingBatchGenerator(schema)41        ... finished = False42        ... while not finished:43        ...     batch = gen(input_seq_len, output_seq_len)44        ...     try:45        ...         training_step(model, batch)46        ...         oom = False47        ...     except torch.cuda.OutOfMemoryError:48        ...         oom = True49        ...     finished = gen.advance(oom)50        ... solution = gen.max_batch_size  # The solution of the search problem.51        ... gen.reset()  # Can re-use for other sequence lengths now.52 53    The search terminates once the difference between max working batch size and min OOM batch size54    divided by the latter is smaller than ``rel_gap_thresh`` that difference amounts to a single element.55    For example, a max working batch size is 96 and min OOM batch size is 100 indicates a gap of 0.04,56    which would terminate the search with threshold of 0.05.57 58    In order to generate mini-batches compatible with a given model, the generator:59 60    * accepts a ``schema`` argument in its constructor, and61 62    * accepts input/output sequence lengths in each call to generate a mini-batch.63 64    ``schema`` has the following structure::65 66 67        >>> {68        ...     "cls": tuple | MyBatchType,69        ...     "inputs": [70        ...         {71        ...             "type": NeuralType(...) | Literal["dummy"],72        ...             "seq_length": Literal["input", "output"],73        ...             "vocab_size": int,  # optional, required only for LabelsType74        ...             "name": str,  # optional, indicates kwarg75        ...         },76        ...         ...,77        ...     ]78        ... }79 80    ``cls`` indicates how we should construct the mini-batch. Typically you can just use ``tuple`` for most81    batch schemas. However, if the model expects a specific, e.g., dataclass, you can tell ``ProfilingBatchGenerator``82    to use it. The mini-batch object will be constructed using the items in ``inputs``.83 84    Each element of ``inputs`` specifies a NeMo NeuralType which needs to have a defined ``elements_type``.85    The supported types are ``AudioSignal``, ``LengthsType`` and ``LabelsType``.86    If "type" is not a NeuralType, we interpret that as a placeholder tensor that's not relevant but expected87    by the model/batch constructor. In addition, ``"seq_length"`` key is used to determine whether we should apply88    input or output sequence length to a given tensor.89 90    Optional keys:91 92    * ``vocab_size`` is required for ``LabelsType`` so that we can generate proper label values.93 94    * ``name`` is required if objects of ``cls`` have to be constructed using keyword arguments.95 96    A simple schema example for a model using audio/lengths tensor pair (unsupervised/self-supervised)::97 98        >>> {99        ...     "cls": tuple,100        ...     "inputs": [101        ...         {"type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input"},102        ...         {"type": NeuralType(("B"), LengthsType()), "seq_length": "input"},103        ...     ]104        ... }105 106    """107 108    def __init__(109        self,110        schema: dict,111        start_batch_size: int = 32,112        rel_gap_thresh: float = 0.05,113        device: str = "cuda",114    ):115        self.schema = schema116        self.start_batch_size = start_batch_size117        self.rel_gap_thresh = rel_gap_thresh118        self.device = device119        self.reset()120 121    def __call__(self, input_seq_length: int, output_seq_length: int):122        B = self._current123        select_seq_length = {"input": input_seq_length, "output": output_seq_length}124        batch = []125        names = []126        for item in self.schema["inputs"]:127            nt = item["type"]128            if isinstance(nt, str) and nt == "constant":129                if isinstance(val := item["value"], str) and val == "batch":130                    tnsr = torch.tensor([B], dtype=torch.long, device=self.device)131                else:132                    tnsr = torch.tensor([val], dtype=torch.long, device=self.device)133            elif not isinstance(nt, NeuralType):  # placeholder134                tnsr = torch.tensor([])135            elif isinstance(nt.elements_type, AudioSignal):136                seq_length = select_seq_length[item["seq_length"]]137                tnsr = torch.randn(B, seq_length, dtype=torch.float32, device=self.device)138            elif isinstance(nt.elements_type, LengthsType):139                seq_length = select_seq_length[item["seq_length"]]140                tnsr = torch.ones(B, dtype=torch.long, device=self.device) * seq_length141            elif isinstance(nt.elements_type, LabelsType):142                seq_length = select_seq_length[item["seq_length"]]143                tnsr = torch.randint(0, item["vocab_size"], size=(B, seq_length), device=self.device)144            elif isinstance(nt.elements_type, MaskType):145                seq_length = select_seq_length[item["seq_length"]]146                tnsr = torch.ones(B, seq_length, device=self.device)147            else:148                raise RuntimeError("Unexpected item in oomptimizer schema: {item}")149            batch.append(tnsr)150            names.append(item.get("name"))151        args = [elem for name, elem in zip(names, batch) if name is None]152        kwargs = {name: elem for name, elem in zip(names, batch) if name is not None}153        if not kwargs and self.schema["cls"] == tuple:154            return tuple(args)155        return self.schema["cls"](*args, **kwargs)156 157    @property158    def max_batch_size(self) -> int | None:159        """160        Return the solution of the batch size search problem.161        It will keep returning None until the search is done.162        """163        if (164            self._max_ok is not None165            and self._min_err is not None166            and (self.current_rel_gap <= self.rel_gap_thresh or self._min_err - self._max_ok <= 1)167        ):168            return self._max_ok169        return None170 171    @property172    def current_rel_gap(self) -> float | None:173        """174        Return the current gap between the largest batch that works and the smallest batch that triggers OOM.175        The gap is defined as the batch size difference divided by the larger element.176        E.g., if the best found batch size is 95 and the smallest that triggers OOM is 100, the gap is 0.05.177        """178        if self._min_err is None or self._max_ok is None:179            return None180        return (self._min_err - self._max_ok) / self._min_err181 182    def reset(self):183        """Reset the generator to prepare it for a new search."""184        self._current = self.start_batch_size185        self._max_ok = None  # max batch size that works186        self._min_err = None  # min batch size that doesn't work187 188    def advance(self, oom: bool) -> bool:189        """190        Adjusts the current batch size based on the outcome.191        Returns a bool indicating whether the calibration is complete.192        """193        if self.max_batch_size is not None:194            return True195 196        if oom:197            # Training step failed with OOM.198            # Update the minimum known batch size that causes an error.199            self._min_err = min(float("inf") if self._min_err is None else self._min_err, self._current)200            # Training step failed on OOM201            if self._max_ok is None:202                # We haven't found a batch size that works yet, keep going 2x down.203                self._current = round(self._current / 2)204            else:205                # Try the middle-point between the known extremes.206                self._current = round((self._max_ok + self._min_err) / 2)207        else:208            # Training step successful.209            # Update the maximum known batch size that works.210            self._max_ok = max(-1 if self._max_ok is None else self._max_ok, self._current)211            if self._min_err is None:212                # We haven't found a batch size that causes an error yet, keep going 2x higher213                self._current *= 2214            else:215                # Try the middle-point between the known extremes.216                self._current = round((self._max_ok + self._min_err) / 2)217 218        return False219 220 221class FloatList(click.Option):222    """Support passing bucket duration bins as [1.1,2.5,5.6,...]"""223 224    name = "list[float]"225 226    def type_cast_value(self, ctx, value):227        if isinstance(value, list) and all(isinstance(v, float) for v in value):228            return value229        try:230            import ast231 232            ans = ast.literal_eval(value)233            if isinstance(ans[0], list):234                ans = [tuple(item) for item in ans]235            return ans236        except ValueError:237            raise click.BadParameter(value)238 239 240@click.command(context_settings={'show_default': True})241@click.option(242    "-n",243    "--pretrained-name",244    type=str,245    default=None,246    help="Name of a pretrained model to use, e.g. 'nvidia/canary-1b'.",247)248@click.option(249    "-m",250    "--module-name",251    type=str,252    default=None,253    help="Full path to NeMo's module corresponding to CONFIG_PATH, e.g. 'nemo.collections.asr.models.EncDecMultiTaskModel'.",254)255@click.option(256    "-c", "--config-path", type=str, default=None, help="Path to the training configuration file for MODULE_NAME."257)258@click.option("-o", "--optimizer-name", type=str, default="adamw", help="Name of optimizer to use.")259@click.option(260    "-b",261    "--buckets",262    cls=FloatList,263    default=[5.0, 10.0, 15.0, 20.0, 25.0, 30.0],264    help="List of upper-bound bucket bins (i.e. first bucket is [0.0 - item0), second bucket is [item0 - item1), etc.). "265    "We also support a nested list for 2D bucketing, e.g. [[2.0, 10],[2.0,20],[4.5,15],[4.5,30],...], "266    "where each item is a pair of (max_input_seq_len, max_output_seq_len) for a given bucket.",267)268@click.option(269    "-t",270    "--threshold",271    type=float,272    default=0.05,273    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.",274)275@click.option("-s", "--start-batch-size", type=int, default=32, help="Initial batch size to start the search from.")276@click.option(277    "-r",278    "--ratio",279    type=int,280    default=12,  # conservative estimate towards longer transcripts281    help="The output_sequence_length to input_sequence_length ratio for the purpose of determing the maximum output sequence lengths. "282    "The interpretation depends on input and output modalities. Examples: for audio->text it's tokens per second. "283    "For text->audio it's seconds per token. For audio->audio it's output seconds per input second. "284    "For text->text it's output tokens per input token. "285    "In general larger ratio means longer output sequences and increased memory consumption. "286    "The default value is set adequately for automatic speech recognition. "287    "This argument is ignored when 2D buckets are provided to --buckets option.",288)289@click.option(290    "-f",291    "--memory-fraction",292    type=float,293    default=0.9,294    help="Limits the use of CUDA memory for this process to MEMORY_FRACTION of the total device memory. "295    "By default we force 5% memory to be unused to account for non-training-loop related CUDA memory usage"296    "in actual training scripts.",297)298@click.option(299    "-d",300    "--device",301    default="cuda:0",302    help="Device string to be passed to torch.device; due to MEMORY_FRACTION option, "303    "it must specify the device index (e.g. cuda:0). "304    "You can also leave the default index and select a specific GPU using env var CUDA_VISIBLE_DEVICES=<idx>",305)306@click.option(307    "-y",308    "--dtype",309    default="bfloat16",310    help="Float precision to use for computation (used together with autocast).",311)312@click.option(313    "--ddp/--no-ddp",314    type=bool,315    default=True,316    help="Whether we should simulate DDP GPU RAM usage. Stores an extra copy of the model in GPU memory. Enabled by default.",317)318def oomptimizer(319    pretrained_name: str | None,320    module_name: str | None,321    config_path: str | None,322    optimizer_name: str,323    buckets: list[float],324    threshold: float,325    start_batch_size: int,326    ratio: int,327    memory_fraction: float,328    device: str,329    dtype: str,330    ddp: bool,331):332    """333    OOMptimizer finds the optimal batch sizes for training your model with bucketing dataloading.334    It performs a search over batch sizes until it converges by measuring the GPU memory usage for335    a model's training step and optimizer update.336 337    \b338    There are two main usage patterns: for using a pretrained model or an untrained model configuration.339    The latter is more flexible but requires the user to provide two separate arguments. Examples:340    * python oomptimizer.py --pretrained-name nvidia/canary-1b341    * python oomptimizer.py --module-name nemo.collections.asr.models.EncDecMultiTaskModel \342        --config-path examples/asr/conf/speech_multitask/fast-conformer_aed.yaml343 344    Dynamic bucketing is notoriously difficult to tune as you risk running into CUDA OOM many steps into the training.345    In order to simplify finding the optimal settings, OOMptimizer scans each bucket to find the maximum possible346    batch size that doesn't trigger a CUDA OOM.347 348    \b349    The suggested workflow is the following:350    1) Run scripts/speech_recognition/estimate_duration_bins.py to get the duration distribution of your data.351        (consider running estimate_duration_bins_2d.py for models with a strong dependency on output sequence length352        such as attention-encoder-decoder models).353    2) Run OOMptimizer to find the optimal batch sizes for your specific model, optimizer, and GPU.354    3) Use these optimal settings in your actual training script and enjoy optimal GPU utilization OOM-free.355 356    In the unlikely event that OOMptimizer bucket batch sizes are still leading to OOMs,357    please try a lower setting of the MEMORY_FRACTION option, e.g. 0.75 (75% of GPU memory).358    This may be required in very complex setups where there are additional GPU RAM loads that can't be anticipated359    through the combination of training_step and optimizer update.360    """361    if all(opt is None for opt in (pretrained_name, module_name, config_path)):362        click.secho(363            "You need to provide either PRETRAINED_NAME or the pair of MODULE_NAME and CONFIG_PATH.", fg="yellow"364        )365        sys.exit(1)366    logging.setLevel(logging.CRITICAL)367    torch.cuda.set_per_process_memory_fraction(memory_fraction, device)368 369    trainer = pl.Trainer(barebones=True)370    trainer.log_every_n_steps = 1000000371    model_clones = []372    for _ in range(2 if ddp else 1):373        if pretrained_name is not None:374            assert (375                config_path is None and module_name is None376            ), "--pretrained-name cannot be used together with --module-name/--config-path"377            click.echo(f"Intializing ASR model from pretrained checkpoint {pretrained_name}.")378            model = ASRModel.from_pretrained(pretrained_name, trainer=trainer).to(device)379        else:380            assert config_path is not None, "--module-name requires --config-path to be specified as well."381            assert module_name is not None, "--config-path requires --module-name to be specified as well."382            cfg = OmegaConf.load(config_path)383            namespace, name = module_name.rsplit('.', maxsplit=1)384            model_cls = getattr(importlib.import_module(namespace), name)385            model = model_cls(cfg=cfg.model, trainer=trainer).to(device)386        model_clones.append(model)387    model = model_clones[-1]388 389    if not hasattr(model, "oomptimizer_schema"):390        click.secho(391            f"We read model of type {type(model)} which doesn't seem to support OOMptimizer "392            f"(we could not find the property .oomptimizer_schema).",393            fg="red",394        )395        sys.exit(1)396 397    schema = model.oomptimizer_schema398 399    click.echo("Setting up the optimizers.")400    optimizer, _ = model.setup_optimization({"name": optimizer_name, "lr": 1e-7, "weight_decay": 0.0})401 402    is_2d_bucketing = all(403        isinstance(item, (list, tuple)) and len(item) == 2 and all(isinstance(v, Number) for v in item)404        for item in buckets405    )406    # Determine modality for input and output.407    modalities = [408        (409            "text"410            if any(411                isinstance(item["type"], NeuralType)412                and isinstance(item["type"].elements_type, LabelsType)413                and item["seq_length"] == direction414                for item in schema["inputs"]415                if item["type"] != "dummy"416            )417            else "audio"418        )419        for direction in ("input", "output")420    ]421 422    def get_max_seq_lens(buckets):423 424        def _determine_lens_for_bucket(bin):425            if is_2d_bucketing:426                input_len, output_len = bin427            else:428                input_len = bin429                output_len = math.ceil(ratio * input_len)430            sampling_rate = getattr(431                model, "sample_rate", 16000432            )  # TODO: may need to extend schema for broader model coverage433            match modalities:434                case "audio", "audio":435                    return (436                        compute_num_samples(input_len, sampling_rate=sampling_rate),437                        compute_num_samples(output_len, sampling_rate=sampling_rate),438                    )439                case "audio", "text":440                    return (compute_num_samples(input_len, sampling_rate=sampling_rate), output_len)441                case "text", "audio":442                    return (443                        input_len,444                        compute_num_samples(output_len, sampling_rate=sampling_rate),445                    )446                case "text", "text":447                    return input_len, output_len448                case _:449                    raise RuntimeError(f"Unexpected modality combination: {_}")450 451        return [_determine_lens_for_bucket(bin) for bin in buckets]452 453    click.echo("Starting profiling.")454    max_seq_lens = get_max_seq_lens(buckets)455    gen = ProfilingBatchGenerator(schema=schema, start_batch_size=start_batch_size, rel_gap_thresh=threshold)456    profile = {}457 458    # Iterate buckets from the largest to the smallest sequences. This usually ends up creating459    # a tiny bit smaller batches, likely due to worse memory fragmentation.460    with torch.autocast("cuda", getattr(torch, dtype)):461        for bucket, (seq_len_in, seq_len_out) in reversed(list(zip(buckets, max_seq_lens))):462            click.echo(f"The current sequence lengths are: input={seq_len_in} output={seq_len_out}.")463            gen.reset()464            batch_idx = 0465 466            def step():467                click.echo(468                    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]"469                )470                batch = gen(seq_len_in, seq_len_out)471                oom = False472                try:473                    click.echo(f"\tCurrent gap: {gen.current_rel_gap}... ", nl=False)474                    optimizer.zero_grad()475                    out = model.training_step(batch, batch_idx)476                    out['loss'].sum().backward()477                    optimizer.step()478                except torch.cuda.OutOfMemoryError as e:479                    click.secho(f"OOM!", fg="yellow")480                    oom = True481                except RuntimeError as e:482                    if "cuFFT error: CUFFT_INTERNAL_ERROR" not in str(e):483                        raise484                    click.secho(f"OOM!", fg="yellow")485                    oom = True486                else:487                    click.secho(f"OK!", fg="green")488                finally:489                    click.echo(490                        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]"491                    )492                    del batch493                    # Note: We could call empty_cache() to free up some more memory on the GPU,494                    #       but we have found out empirically that this causes a mismatched condition495                    #       between OOMptimizer and the actual training. During training, there is some496                    #       degree of memory fragmentation and it's better to simulate that in OOMptimizer.497                    # torch.cuda.memory.empty_cache()498                    torch.cuda.reset_max_memory_allocated()499                return oom500 501            oom = step()502            while not (finished := gen.advance(oom)):503                click.echo("\t" + "=" * 80)504                oom = step()505 506            click.secho(507                f"=> Optimal setting for bucket={bucket} (input={seq_len_in} output={seq_len_out}) is max_batch_size={gen.max_batch_size}",508                fg="green",509            )510            profile[(bucket, seq_len_in, seq_len_out)] = gen.max_batch_size511            gen.start_batch_size = gen.max_batch_size * 2512 513    # Reverse the profile to be ascendingly sorted again.514    profile = dict(reversed(list(profile.items())))515 516    click.echo("The 1st stage profile is:")517    for (bucket, seq_len_in, seq_len_out), bs in profile.items():518        click.echo(f"Bucket={bucket} (input={seq_len_in} output={seq_len_out}) => max_batch_size={bs}")519 520    if is_2d_bucketing:521        # 2D bucketing doesn't support bucket merging.522        final_profile = [["[" + ",".join(map(str, b)) + "]", bs] for (b, _, __), bs in profile.items()]523    else:524        click.echo("Bucket merging stage...")525        final_profile = []526        for idx, ((bucket, seq_len_in, seq_len_out), bs) in enumerate(profile.items()):527            if idx == 0:528                final_profile.append([bucket, bs])529                continue530            if bs == final_profile[-1][1]:531                click.echo(f"Merging bucket {idx} with bucket {idx-1} due to identical batch sizes.")532                final_profile[-1][0] = bucket533                continue534            final_profile.append([bucket, bs])535 536    click.secho(f"The profile was created with the following settings:")537    click.secho(f"* using {memory_fraction:.1%} of available GPU RAM.")538    click.secho(f"* {'' if ddp else 'not '}simulating DDP memory overhead.")539    click.secho(f"* using AMP with dtype={dtype}.")540    click.secho("The final profile is:", bold=True)541    click.secho("\tbucket_duration_bins=[" + ",".join(str(seqlen) for seqlen, bs in final_profile) + "]", bold=True)542    click.secho("\tbucket_batch_size=[" + ",".join(str(bs) for seqlen, bs in final_profile) + "]", bold=True)543 544 545if __name__ == "__main__":546    oomptimizer()547