Gilkeren/debug2

src/fairseq2/assets/cards/datasets/librilight.yaml

@@ -7,3 +7,15 @@
 name: librilight_asr_10h
 dataset_family: generic_asr
 data: /checkpoint/mms/shared/assets/debug/librilight
+
+---
+
+name: librilight_asr_10h_unlabeled
+dataset_family: generic_speech
+data: /checkpoint/mms/shared/assets/debug/librilight
+
+---
+
+name: librilight_asr_10h_inference
+dataset_family: speech_inference
+data: /checkpoint/mms/shared/assets/debug/librilight

src/fairseq2/cli/_main.py

@@ -8,47 +8,47 @@
 
 import os
 import sys
-from signal import SIG_DFL, SIGINT, raise_signal, signal
+from signal import raise_signal, SIG_DFL, SIGINT, signal
 
 import torch
-from torch.cuda import OutOfMemoryError
 
 from fairseq2 import setup_fairseq2
-from fairseq2.cli.utils.rich import create_rich_progress_reporter
-from fairseq2.error import ContractError, InternalError
-from fairseq2.extensions import ExtensionError
-from fairseq2.logging import LoggingSetupError, log
-from fairseq2.setup import SetupError
-from fairseq2.utils.env import InvalidEnvironmentVariableError, get_rank
 
 # isort: split
 
 from fairseq2.cli._logging import setup_logging
 from fairseq2.cli._setup import setup_cli
+from fairseq2.cli.utils.rich import create_rich_progress_reporter
+from fairseq2.error import ContractError, InternalError
+from fairseq2.extensions import ExtensionError
+from fairseq2.logging import log, LoggingSetupError
+from fairseq2.setup import SetupError
+from fairseq2.utils.env import get_rank, InvalidEnvironmentVariableError
+from torch.cuda import OutOfMemoryError
 
 
 def main() -> None:
     """Runs the command line fairseq2 program."""
     exit_code = 1
 
-    try:
-        exit_code = _run()
-    except KeyboardInterrupt:
-        log.info("Command canceled!")
+    # try:
+    exit_code = _run()
+    # except KeyboardInterrupt:
+    #     log.info("Command canceled!")
 
-        signal(SIGINT, SIG_DFL)
+    #     signal(SIGINT, SIG_DFL)
 
-        raise_signal(SIGINT)
-    except OutOfMemoryError:
-        s = torch.cuda.memory_summary()
+    #     raise_signal(SIGINT)
+    # except OutOfMemoryError:
+    #     s = torch.cuda.memory_summary()
 
-        log.exception("CUDA out of memory. See logged memory stats.\n{}", s)
-    except InternalError:
-        log.exception("Command failed with an unexpected internal error. Please file a bug report.")  # fmt: skip
-    except ContractError:
-        log.exception("Command failed with an unexpected internal error caused by an extension. Please file a bug report to the corresponding extension author.")  # fmt: skip
-    except Exception:
-        log.exception("Command failed with an unexpected error. See the logged stack trace for details.")  # fmt: skip
+    #     log.exception("CUDA out of memory. See logged memory stats.\n{}", s)
+    # except InternalError:
+    #     log.exception("Command failed with an unexpected internal error. Please file a bug report.")  # fmt: skip
+    # except ContractError:
+    #     log.exception("Command failed with an unexpected internal error caused by an extension. Please file a bug report to the corresponding extension author.")  # fmt: skip
+    # except Exception:
+    #     log.exception("Command failed with an unexpected error. See the logged stack trace for details.")  # fmt: skip
 
     sys.exit(exit_code)
 
@@ -84,3 +84,7 @@ def _run() -> int:
         return 1
 
     return cli.run(context)
+
+
+if __name__ == "__main__":
+    main()

src/fairseq2/models/llama/_config.py

@@ -224,6 +224,25 @@ def llama3_1_8b() -> LLaMAConfig:
 
         return config
 
+    @arch("llama3_1_8b_v4_tokenizer")
+    def llama3_1_8b_v4_tokenizer() -> LLaMAConfig:
+        config = llama3_1_8b()
+        config.vocab_size = 9812
+        config.pad_idx = 1
+        config.model_dim = 2048
+        config.tie_embeddings = True  # remapped from tied_embeddings
+        config.ffn_inner_dim = 2048 * 4
+        config.ffn_inner_dim_multiplier = 1.5
+        config.ffn_inner_dim_to_multiple = (
+            256  # remapped from ffn_inner_dim_multiple_of
+        )
+        config.num_attn_heads = 32
+        config.num_key_value_heads = 8
+        config.num_layers = 16
+        config.use_scaled_rope = True
+        config.rope_scaling.factor = 32.0  # renamed from rope_scale.factor
+        return config
+
     @arch("llama3_1_70b")
     def llama3_1_70b() -> LLaMAConfig:
         config = llama3_70b()

src/fairseq2/models/wav2vec2/asr/_config.py

@@ -500,3 +500,51 @@ def v4_tokenizer_updated_300m() -> Wav2Vec2AsrConfig:
             pad_idx=1,
         )
         return config
+
+    @wav2vec2_asr_arch("7b_v7_tokenizer")
+    def v7_tokenizer_7b() -> Wav2Vec2AsrConfig:
+        config = bib1143_7b()
+        config.vocab_info = VocabularyInfo(
+            size=9818,
+            unk_idx=3,
+            bos_idx=0,
+            eos_idx=2,
+            pad_idx=1,
+        )
+        return config
+
+    @wav2vec2_asr_arch("300m_v7_tokenizer")
+    def v7_tokenizer_300m() -> Wav2Vec2AsrConfig:
+        config = bib1143_300m()
+        config.vocab_info = VocabularyInfo(
+            size=9818,
+            unk_idx=3,
+            bos_idx=0,
+            eos_idx=2,
+            pad_idx=1,
+        )
+        return config
+
+    @wav2vec2_asr_arch("1b_v7_tokenizer")
+    def v7_tokenizer_1b() -> Wav2Vec2AsrConfig:
+        config = bib1143_1b()
+        config.vocab_info = VocabularyInfo(
+            size=9818,
+            unk_idx=3,
+            bos_idx=0,
+            eos_idx=2,
+            pad_idx=1,
+        )
+        return config
+
+    @wav2vec2_asr_arch("3b_v7_tokenizer")
+    def v7_tokenizer_3b() -> Wav2Vec2AsrConfig:
+        config = bib1143_3b()
+        config.vocab_info = VocabularyInfo(
+            size=9818,
+            unk_idx=3,
+            bos_idx=0,
+            eos_idx=2,
+            pad_idx=1,
+        )
+        return config

src/fairseq2/nn/utils/module.py

@@ -13,13 +13,13 @@
 from typing import Protocol, runtime_checkable
 
 import torch
-from torch import Tensor
-from torch.nn import Module, Parameter
-from torch.nn.utils import remove_weight_norm  # type: ignore[attr-defined]
 
 from fairseq2.gang import Gang
 from fairseq2.logging import log
 from fairseq2.typing import CPU, Device
+from torch import Tensor
+from torch.nn import Module, Parameter
+from torch.nn.utils import remove_weight_norm  # type: ignore[attr-defined]
 
 
 @runtime_checkable
@@ -210,7 +210,12 @@ def collect_tensors(m: Module) -> None:
 
     # Do not memoize. No need anyways, and would also break the sync between the
     # traversed tensors and the iterator.
-    apply_to_parameters(target_module, lambda _: next(it), no_memo=True)
+    # apply_to_parameters(
+    #     target_module, lambda _: next(it), no_memo=True, skip_freqs=True
+    # )
+    apply_to_parameters(
+        target_module, lambda _: next(it), no_memo=True, skip_freqs=False
+    )
 
 
 def apply_to_parameters(
@@ -220,6 +225,7 @@ def apply_to_parameters(
     recurse: bool = True,
     memo: dict[Tensor, Tensor] | None = None,
     no_memo: bool = False,
+    skip_freqs: bool = False,
 ) -> None:
     """Apply ``fn`` to the parameters and buffers of ``module``.
 
@@ -245,7 +251,12 @@ def apply_to_parameters(
         for child in module.children():
             if child is not None:
                 apply_to_parameters(
-                    child, fn, recurse=recurse, memo=memo, no_memo=no_memo
+                    child,
+                    fn,
+                    recurse=recurse,
+                    memo=memo,
+                    no_memo=no_memo,
+                    skip_freqs=skip_freqs,
                 )
 
     def call_fn(
@@ -273,6 +284,12 @@ def call_fn(
         with torch.no_grad():
             new_param = call_fn(param, is_param=True, requires_grad=param.requires_grad)
 
+        if param.shape != new_param.shape:
+            log.warning(
+                f"The shape of {param_name} changed from {param.shape} to {new_param.shape}."
+            )
+            continue
+
         setattr(module, param_name, new_param)
 
         if (grad := param.grad) is not None:
@@ -285,6 +302,15 @@ def call_fn(
         if buffer is None:
             continue
 
+        if skip_freqs and buffer_name == "freqs":
+            log.warning(
+                f"The `freqs` buffer of `module` was not updated :{buffer_name}."
+            )
+            target = call_fn(buffer)
+            log.info(f"{target=} {buffer=}")
+            setattr(module, buffer_name, buffer)
+            continue
+
         setattr(module, buffer_name, call_fn(buffer))
 
 
@@ -464,6 +490,21 @@ def load_state_dict(
     ``state_dict`` does not contain any keys corresponding to descendants that are set to ``None``
     via :meth:`Module.register_module()`.
     """
+    # Key mapping
+    need_mapping = False
+    sample_key = list(state_dict.keys())[0]
+    if (
+        sample_key.startswith("module.")
+        and not sample_key in module.state_dict().keys()
+    ):
+        mapped_key = sample_key[7:]
+        if mapped_key in module.state_dict().keys():
+            need_mapping = True
+
+    if need_mapping:
+        key_mapping = lambda key: key[7:] if key.startswith("module.") else key
+        state_dict = {key_mapping(key): value for key, value in state_dict.items()}
+
     module.load_state_dict(state_dict, strict=strict)
 
     unexpected_keys = []

src/fairseq2/recipes/_trainer.py

@@ -9,16 +9,10 @@
 from abc import ABC, abstractmethod
 from contextlib import nullcontext
 from enum import Enum
-from typing import Final, Generic, Mapping, TypeVar, final
+from typing import Final, final, Generic, Mapping, TypeVar
 
 import torch
 import torch.distributed
-from rich.pretty import pretty_repr
-from torch import Tensor
-from torch.cuda import OutOfMemoryError
-from torch.optim import Optimizer
-from torch.profiler import record_function
-from typing_extensions import override
 
 from fairseq2.checkpoint import (
     CheckpointError,
@@ -30,21 +24,14 @@
 from fairseq2.datasets import DataReader, DataReadError
 from fairseq2.device import SupportsDeviceTransfer
 from fairseq2.error import InternalError, InvalidOperationError
-from fairseq2.gang import GangError, Gangs, broadcast_flag
+from fairseq2.gang import broadcast_flag, GangError, Gangs
 from fairseq2.logging import log
 from fairseq2.metrics import Mean, MetricBag, MetricBagError, MetricDescriptor
 from fairseq2.metrics.recorders import MetricRecorder, MetricRecordError
 from fairseq2.nn.utils.gradient import check_gradient_norms, normalize_gradients
 from fairseq2.optim import DynamicLossScaler
-from fairseq2.optim.lr_scheduler import LRScheduler, get_effective_lr
+from fairseq2.optim.lr_scheduler import get_effective_lr, LRScheduler
 from fairseq2.profilers import Profiler
-from fairseq2.typing import CPU, ContextManager, DataType
-from fairseq2.utils.device_stat import DeviceStatTracker
-from fairseq2.utils.gc import GarbageCollector
-from fairseq2.utils.progress import ProgressReporter, ProgressTask
-from fairseq2.utils.rng import RngBag
-from fairseq2.utils.state import Stateful
-from fairseq2.utils.stopwatch import Stopwatch
 
 # isort: split
 
@@ -59,6 +46,19 @@
 from fairseq2.recipes._model import Model
 from fairseq2.recipes._recipe import Recipe, RecipeStopException
 from fairseq2.recipes._validator import Validator
+from fairseq2.typing import ContextManager, CPU, DataType
+from fairseq2.utils.device_stat import DeviceStatTracker
+from fairseq2.utils.gc import GarbageCollector
+from fairseq2.utils.progress import ProgressReporter, ProgressTask
+from fairseq2.utils.rng import RngBag
+from fairseq2.utils.state import Stateful
+from fairseq2.utils.stopwatch import Stopwatch
+from rich.pretty import pretty_repr
+from torch import Tensor
+from torch.cuda import OutOfMemoryError
+from torch.optim import Optimizer
+from torch.profiler import record_function
+from typing_extensions import override
 
 BatchT_contra = TypeVar(
     "BatchT_contra", bound=SupportsDeviceTransfer, contravariant=True
@@ -718,7 +718,12 @@ def _do_run_step(self, progress_task: ProgressTask) -> _TrainerState:
                 batch = batches.pop()
 
                 try:
-                    batch.to(gangs.root.device)
+                    try:
+                        batch.to(gangs.root.device)
+                    except Exception as e:
+                        log.info(f"{gangs.root.device=}")
+                        log.info(f"{batch=}")
+                        raise e
 
                     with self._maybe_no_sync(batch_nr, num_batches):
                         with record_function(f"step_{step_nr}_{batch_nr}_forward"):