RASMUS/Finnish-ASR-Canary-v2

1# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import json16import logging17import os.path18from io import BytesIO19from pathlib import Path20from typing import Any, Dict, Union21 22import numpy23 24# tenosrstore is needed to register 'bfloat16' dtype with numpy for zarr compatibility25import tensorstore  # noqa: F401 pylint: disable=unused-import26import torch27from torch.distributed.checkpoint import FileSystemReader, load28from torch.distributed.checkpoint.metadata import BytesStorageMetadata, TensorStorageMetadata29 30from nemo.export.tarutils import TarPath, ZarrPathStore31from nemo.export.utils._mock_import import _mock_import32 33LOGGER = logging.getLogger("NeMo")34 35 36def nemo_to_path(nemo_checkpoint: Union[Path, str]) -> Union[Path, TarPath]:37    """38    Creates Path / TarPath object suitable for navigating inside the nemo checkpoint.39 40    Args:41        nemo_checkpoint (Path, str): Path to the NeMo checkpoint.42    Returns:43        Path | TarPath: Suitable Path object for navigating through the checkpoint.44    """45    string_path = str(nemo_checkpoint)46 47    if os.path.isdir(string_path):48        return Path(string_path)49    return TarPath(string_path)50 51 52class TarFileSystemReader(FileSystemReader):53    """Reader that accepts both Path and TarPath checkpoint directory.54 55    The FileSystemReader works with TarPath, but expects a pure Path.56    It's enough to skip the Path check in __init__.57    """58 59    def __init__(self, path: Union[Path, TarPath]) -> None:60        """Makes sure that super().__init__ gets a pure path as expected."""61        super_path = str(path) if isinstance(path, TarPath) else path62        super().__init__(super_path)63        if isinstance(path, TarPath):64            self.path = path  # overwrites path set in super().__init__ call65 66 67def load_sharded_metadata_torch_dist(68    checkpoint_dir: Union[Path, TarPath], load_extra_states: bool = False69) -> Dict[str, Any]:70    """71    Loads model state dictionary from torch_dist checkpoint.72 73    Args:74        checkpoint_dir (Path | TarPath): Path to the model weights directory.75        load_extra_states (bool): If set to true, loads BytesIO objects, related to the extra states.76    Returns:77        dict: Loaded model state dictionary (weights are stored in torch tensors).78    """79    fs_reader = TarFileSystemReader(checkpoint_dir)80    metadata = fs_reader.read_metadata()81 82    state_dict = {83        k: torch.empty(tp.size, dtype=tp.properties.dtype)84        for k, tp in metadata.state_dict_metadata.items()85        if isinstance(tp, TensorStorageMetadata)86    }87 88    if load_extra_states:89        state_dict.update(90            {k: [] for k, tp in metadata.state_dict_metadata.items() if isinstance(tp, BytesStorageMetadata)}91        )92 93    load(state_dict, storage_reader=fs_reader)94    return state_dict95 96 97def load_sharded_pickle_extra_state_scale(dir: Union[Path, TarPath]) -> Dict[str, BytesIO]:98    """99    Loads model extra states from the .pt shards.100 101    Args:102        dir (Path | TarPath): Path to the directory with sharded extra states.103    Returns:104        dict: State dictionary corresponding to the loaded extra states.105    """106    pt_files = list(dir.glob('shard_*_*.pt'))107    extra_states = {}108    for file in pt_files:109        shard_name = file.name.split('.')[0]110        with file.open('rb') as opened_file:111            extra_states[dir.name + '/' + shard_name] = torch.load(opened_file, weights_only=True)112 113    return extra_states114 115 116def contains_extra_states(subdir: Union[Path, TarPath]) -> bool:117    """118    Checks if zarr directory contains extra states.119 120    Args:121        subdir (Path | TarPath): Directory inside the zarr checkpoint.122    Returns:123        bool: Is a directory with extra states124    """125    return list(subdir.glob('shard_0_*.pt')) != []126 127 128def load_sharded_metadata_zarr(129    checkpoint_dir: Union[Path, TarPath], load_extra_states: bool = False130) -> Dict[str, Any]:131    """132    Loads model dictionary from the zarr format.133 134    Args:135        checkpoint_dir (Path | TarPath): Path to the NeMo checkpoint.136        load_extra_states (bool): If set to True, the function will load BufferIO objects with extra states.137    Returns:138        dict: Model state dictionary.139    """140    if load_extra_states:141        torch.serialization.add_safe_globals([BytesIO])142 143    sharded_state_dict = {}144    for subdir in checkpoint_dir.iterdir():145        if not subdir.is_dir():146            continue147 148        if load_extra_states and contains_extra_states(subdir):149            sharded_state_dict.update(load_sharded_pickle_extra_state_scale(subdir))150 151        elif (subdir / '.zarray').exists():152            key = subdir.name153            zstore = ZarrPathStore(subdir)154 155            import zarr156 157            arr = zarr.open(zstore, 'r')158 159            if arr.dtype.name == "bfloat16":160                sharded_state_dict[key] = torch.from_numpy(arr[:].view(numpy.int16)).view(torch.bfloat16)161            else:162                sharded_state_dict[key] = torch.from_numpy(arr[:])163 164    return sharded_state_dict165 166 167def nemo_weights_directory(nemo_path: Union[Path, TarPath]) -> Union[Path, TarPath]:168    """169    Returns a Path pointing to the weights directory inside the NeMo checkpoint.170 171    Args:172        nemo_path (Path | TarPath): Path to the nemo checkpoint.173    Returns:174        Path | TarPath: Path to the weights directory inside the model checkpoint.175    """176    if (nemo_path / "model_weights").exists():177        return nemo_path / "model_weights"178 179    if (nemo_path / "weights").exists():180        return nemo_path / "weights"181 182    return nemo_path183 184 185def load_model_weights(checkpoint_path: Union[str, Path], load_extra_states: bool = False) -> Dict[str, Any]:186    """187    Loads NeMo state dictionary. Weights are stored in torch.Tensor188 189    Args:190        checkpoint_path (str | Path): Path to the NeMo checkpoint.191        load_extra_states (bool): If True, loads BytesIO objects, corresponding to the extra states.192    Returns:193        dict: Model state dictionary.194    """195 196    nemo_path = nemo_to_path(checkpoint_path)197    nemo_weights = nemo_weights_directory(nemo_path)198 199    with (nemo_weights / 'metadata.json').open(mode='r') as f:200        config_dict = json.load(f)201 202    if config_dict['sharded_backend'] == 'zarr':203        return load_sharded_metadata_zarr(nemo_weights, load_extra_states=load_extra_states)204    elif config_dict['sharded_backend'] == 'torch_dist':205        # TODO: Remove mocking imports once MCore is available in NIM containers206        with _mock_import("megatron.core.dist_checkpointing.strategies.torch"):207            return load_sharded_metadata_torch_dist(nemo_weights, load_extra_states=load_extra_states)208 209    raise NotImplementedError(f'Distributed checkpoint backend {config_dict["sharded_backend"]} not supported')210