nvidia-nemo--speech/tests/collections/speechlm2/test_init_from_checkpoint.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for init_from_checkpoint functionality in speechlm2.

Unit tests use simple nn.Module subclasses (no HF downloads, no CUDA).
Integration tests use real SALM / SALMAutomodel (require HF config download;
SALMAutomodel tests also require CUDA).
"""
import os
from unittest.mock import patch

import pytest
import torch
from omegaconf import DictConfig
from safetensors.torch import save_file

from nemo.collections.speechlm2.parts.pretrained import (
    _is_dcp_checkpoint,
    init_from_training_checkpoint,
    maybe_load_pretrained_models,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


class SimpleModel(torch.nn.Module):
    """Tiny model used for fast, self-contained unit tests."""

    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 4, bias=False)
        self.norm = torch.nn.LayerNorm(4)


class ConfigurableModel(torch.nn.Module):
    """SimpleModel that carries a ``cfg`` attribute, like SALM/SALMAutomodel."""

    def __init__(self, cfg: dict):
        super().__init__()
        self.cfg = DictConfig(cfg)
        self.linear = torch.nn.Linear(4, 4, bias=False)
        self.norm = torch.nn.LayerNorm(4)


def _save_ckpt(model, path):
    """Save model state_dict in Lightning checkpoint format."""
    torch.save({"state_dict": model.state_dict()}, path)


def _assert_state_dicts_equal(sd1, sd2):
    assert set(sd1.keys()) == set(sd2.keys()), f"Key mismatch: {set(sd1.keys()) ^ set(sd2.keys())}"
    for key in sd1:
        assert torch.equal(sd1[key].cpu(), sd2[key].cpu()), f"Tensor mismatch at key: {key}"


def _assert_state_dicts_not_equal(sd1, sd2):
    """Assert at least one tensor differs (sanity-check before loading)."""
    assert set(sd1.keys()) == set(sd2.keys())
    any_diff = any(not torch.equal(sd1[k].cpu(), sd2[k].cpu()) for k in sd1)
    assert any_diff, "Expected state dicts to differ before checkpoint loading"


# ---------------------------------------------------------------------------
# _is_dcp_checkpoint
# ---------------------------------------------------------------------------


class TestIsDcpCheckpoint:
    def test_with_metadata(self, tmp_path):
        ckpt_dir = tmp_path / "step=100.ckpt"
        ckpt_dir.mkdir()
        (ckpt_dir / ".metadata").touch()
        assert _is_dcp_checkpoint(str(ckpt_dir))

    def test_without_metadata(self, tmp_path):
        ckpt_dir = tmp_path / "step=100.ckpt"
        ckpt_dir.mkdir()
        assert not _is_dcp_checkpoint(str(ckpt_dir))

    def test_regular_file(self, tmp_path):
        ckpt_file = tmp_path / "step=100.ckpt"
        ckpt_file.touch()
        assert not _is_dcp_checkpoint(str(ckpt_file))

    def test_nonexistent_path(self, tmp_path):
        assert not _is_dcp_checkpoint(str(tmp_path / "nonexistent"))

    def test_hf_dir_without_metadata(self, tmp_path):
        """HF directory (model.safetensors) should NOT be detected as DCP."""
        hf_dir = tmp_path / "hf_model"
        hf_dir.mkdir()
        (hf_dir / "model.safetensors").touch()
        assert not _is_dcp_checkpoint(str(hf_dir))


# ---------------------------------------------------------------------------
# init_from_training_checkpoint — non-DCP paths
# ---------------------------------------------------------------------------


class TestInitFromTrainingCheckpoint:
    def test_none_is_noop(self):
        model = SimpleModel()
        original = model.linear.weight.clone()
        init_from_training_checkpoint(model, None)
        assert torch.equal(model.linear.weight, original)

    def test_single_file_ckpt(self, tmp_path):
        source = SimpleModel()
        torch.nn.init.ones_(source.linear.weight)

        ckpt_path = str(tmp_path / "model.ckpt")
        _save_ckpt(source, ckpt_path)

        target = SimpleModel()  # different random init
        _assert_state_dicts_not_equal(source.state_dict(), target.state_dict())

        init_from_training_checkpoint(target, ckpt_path)
        _assert_state_dicts_equal(target.state_dict(), source.state_dict())

    def test_hf_directory(self, tmp_path):
        source = SimpleModel()
        torch.nn.init.ones_(source.linear.weight)

        hf_dir = tmp_path / "hf_model"
        hf_dir.mkdir()
        save_file(source.state_dict(), str(hf_dir / "model.safetensors"))

        target = SimpleModel()
        _assert_state_dicts_not_equal(source.state_dict(), target.state_dict())

        init_from_training_checkpoint(target, str(hf_dir))
        _assert_state_dicts_equal(target.state_dict(), source.state_dict())


# ---------------------------------------------------------------------------
# init_from_training_checkpoint — DCP path (mocked)
# ---------------------------------------------------------------------------


class TestInitFromTrainingCheckpointDCP:
    def test_dcp_calls_distributed_load(self, tmp_path):
        """Verify DCP checkpoint triggers torch.distributed.checkpoint.load."""
        ckpt_dir = tmp_path / "step=100.ckpt"
        ckpt_dir.mkdir()
        (ckpt_dir / ".metadata").touch()

        model = SimpleModel()

        with patch("nemo.collections.speechlm2.parts.pretrained.torch.distributed.checkpoint.load") as mock_load:
            init_from_training_checkpoint(model, str(ckpt_dir))

            mock_load.assert_called_once()
            args, kwargs = mock_load.call_args
            state_dict_wrapper = args[0]
            assert "state_dict" in state_dict_wrapper
            assert kwargs["checkpoint_id"] == str(ckpt_dir)

    def test_dcp_state_dict_has_model_keys(self, tmp_path):
        """The state dict passed to dcp.load should contain model parameter keys."""
        ckpt_dir = tmp_path / "step=100.ckpt"
        ckpt_dir.mkdir()
        (ckpt_dir / ".metadata").touch()

        model = SimpleModel()

        with patch("nemo.collections.speechlm2.parts.pretrained.torch.distributed.checkpoint.load") as mock_load:
            init_from_training_checkpoint(model, str(ckpt_dir))

            state_dict_wrapper = mock_load.call_args[0][0]
            model_sd = state_dict_wrapper["state_dict"]
            assert "linear.weight" in model_sd
            assert "norm.weight" in model_sd


# ---------------------------------------------------------------------------
# maybe_load_pretrained_models — init_from_checkpoint config key
# ---------------------------------------------------------------------------


class TestMaybeLoadPretrainedModels:
    def test_init_from_checkpoint_loads_weights(self, tmp_path):
        source = ConfigurableModel(cfg={})
        torch.nn.init.ones_(source.linear.weight)
        ckpt_path = str(tmp_path / "model.ckpt")
        _save_ckpt(source, ckpt_path)

        target = ConfigurableModel(cfg={"init_from_checkpoint": ckpt_path})
        _assert_state_dicts_not_equal(source.state_dict(), target.state_dict())

        maybe_load_pretrained_models(target)
        _assert_state_dicts_equal(target.state_dict(), source.state_dict())

    def test_init_from_checkpoint_null_is_noop(self):
        model = ConfigurableModel(cfg={"init_from_checkpoint": None})
        original = model.linear.weight.clone()
        maybe_load_pretrained_models(model)
        assert torch.equal(model.linear.weight, original)

    def test_init_from_checkpoint_key_missing_is_noop(self):
        model = ConfigurableModel(cfg={})
        original = model.linear.weight.clone()
        maybe_load_pretrained_models(model)
        assert torch.equal(model.linear.weight, original)

    def test_pretrained_s2s_model_still_works(self, tmp_path):
        """Backward compat: pretrained_s2s_model should still be recognized."""
        source = ConfigurableModel(cfg={})
        torch.nn.init.ones_(source.linear.weight)
        ckpt_path = str(tmp_path / "model.ckpt")
        _save_ckpt(source, ckpt_path)

        target = ConfigurableModel(cfg={"pretrained_s2s_model": ckpt_path})
        maybe_load_pretrained_models(target)
        _assert_state_dicts_equal(target.state_dict(), source.state_dict())


# ---------------------------------------------------------------------------
# SALM integration test
# ---------------------------------------------------------------------------

AUDIO_LOCATOR_TAG = "<|audioplaceholder|>"
SALM_PERCEPTION_CFG = {
    "target": "nemo.collections.speechlm2.modules.perception.AudioPerceptionModule",
    "output_dim": 2048,
    "encoder": {
        "_target_": "nemo.collections.asr.modules.ConformerEncoder",
        "att_context_size": [-1, -1],
        "causal_downsampling": False,
        "conv_context_size": None,
        "conv_kernel_size": 9,
        "conv_norm_type": "batch_norm",
        "d_model": 1024,
        "dropout": 0.1,
        "dropout_att": 0.1,
        "dropout_emb": 0.0,
        "dropout_pre_encoder": 0.1,
        "feat_in": 128,
        "feat_out": -1,
        "ff_expansion_factor": 4,
        "n_heads": 8,
        "n_layers": 2,
        "pos_emb_max_len": 5000,
        "self_attention_model": "rel_pos",
        "subsampling": "dw_striding",
        "subsampling_conv_channels": 256,
        "subsampling_factor": 8,
    },
    "modality_adapter": {
        "_target_": "nemo.collections.speechlm2.modules.perception.IdentityConnector",
        "d_model": 1024,
    },
    "preprocessor": {
        "_target_": "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor",
        "dither": 1e-05,
        "features": 128,
        "frame_splicing": 1,
        "log": True,
        "n_fft": 512,
        "normalize": "per_feature",
        "pad_to": 0,
        "pad_value": 0.0,
        "sample_rate": 16000,
        "window": "hann",
        "window_size": 0.025,
        "window_stride": 0.01,
    },
}


def _resolve_pretrained_salm():
    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
        return {
            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/TinyLlama--TinyLlama_v1.1",
            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/canary-1b-flash.nemo",
        }
    return {
        "pretrained_asr": "nvidia/canary-1b-flash",
        "pretrained_llm": "TinyLlama/TinyLlama_v1.1",
    }


def _resolve_pretrained_automodel():
    if os.path.exists("/home/TestData/speechlm/pretrained_models"):
        return {
            "pretrained_llm": "/home/TestData/speechlm/pretrained_models/Qwen--Qwen3-1.7B",
            "pretrained_asr": "/home/TestData/speechlm/pretrained_models/canary-1b-flash.nemo",
        }
    return {
        "pretrained_asr": "nvidia/canary-1b-flash",
        "pretrained_llm": "Qwen/Qwen3-1.7B",
    }


def _make_salm_cfg(**overrides):
    cfg = {
        **_resolve_pretrained_salm(),
        "pretrained_weights": False,
        "prompt_format": "llama2",
        "audio_locator_tag": AUDIO_LOCATOR_TAG,
        "perception": SALM_PERCEPTION_CFG,
        "optimizer": {"_target_": "torch.optim.AdamW"},
    }
    cfg.update(overrides)
    return cfg


def _make_automodel_cfg(**overrides):
    cfg = {
        **_resolve_pretrained_automodel(),
        "pretrained_weights": False,
        "prompt_format": "qwen",
        "audio_locator_tag": AUDIO_LOCATOR_TAG,
        "perception": SALM_PERCEPTION_CFG,
        "optimizer": {"_target_": "torch.optim.AdamW"},
        "torch_dtype": "bfloat16",
    }
    cfg.update(overrides)
    return cfg


def test_salm_init_from_checkpoint(tmp_path):
    from nemo.collections.speechlm2.models import SALM

    # Create source model and save checkpoint
    model1 = SALM(_make_salm_cfg())
    expected_sd = {k: v.clone().cpu() for k, v in model1.state_dict().items()}
    ckpt_path = str(tmp_path / "source.ckpt")
    _save_ckpt(model1, ckpt_path)
    del model1

    # Create target model — init_from_checkpoint overrides the random init
    model2 = SALM(_make_salm_cfg(init_from_checkpoint=ckpt_path))
    _assert_state_dicts_equal(model2.state_dict(), expected_sd)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="SALMAutomodel requires CUDA")
def test_salm_automodel_init_from_checkpoint(tmp_path):
    from nemo.collections.speechlm2.models import SALMAutomodel

    # Create source model and save checkpoint
    model1 = SALMAutomodel(_make_automodel_cfg())
    model1.configure_model()
    expected_sd = {k: v.clone().cpu() for k, v in model1.state_dict().items()}
    ckpt_path = str(tmp_path / "source.ckpt")
    _save_ckpt(model1, ckpt_path)
    del model1
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Create target model — configure_model loads checkpoint via maybe_load_pretrained_models
    model2 = SALMAutomodel(_make_automodel_cfg(init_from_checkpoint=ckpt_path))
    model2.configure_model()
    _assert_state_dicts_equal(model2.state_dict(), expected_sd)