Skip to content

Commit 2c1a0c8

Browse files
committed
tests: FunASRNano opt-ins leave defaults untouched
1 parent d41d1c9 commit 2c1a0c8

1 file changed

Lines changed: 136 additions & 0 deletions

File tree

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
"""llm_conf.sdpa_backends / llm_conf.torch_compile on both copies of FunASRNano (the package
2+
class and the recipe's model.py, imported by path). CPU-only, no weights: tiny stand-ins."""
3+
4+
import importlib.util
5+
import os
6+
import sys
7+
import types
8+
9+
import pytest
10+
import torch
11+
import torch.nn as nn
12+
13+
from funasr.models.fun_asr_nano import model as funasr_model
14+
from funasr.register import tables
15+
16+
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
17+
RECIPE_DIR = os.path.join(REPO, "examples", "industrial_data_pretraining", "fun_asr_nano")
18+
19+
20+
class _TinyDecoder(nn.Module):
21+
"""Stands in for the HF decoder stack (``llm.model``); records every eager call."""
22+
23+
def __init__(self, dim):
24+
super().__init__()
25+
self.proj = nn.Linear(dim, dim)
26+
self.eager_calls = []
27+
28+
def forward(self, input_ids=None, inputs_embeds=None, **kwargs):
29+
if not torch.compiler.is_compiling():
30+
cudnn = torch.backends.cuda.cudnn_sdp_enabled()
31+
self.eager_calls.append({"batch": int(inputs_embeds.shape[0]), "cudnn": cudnn})
32+
return self.proj(inputs_embeds)
33+
34+
35+
class _TinyLLM(nn.Module):
36+
def __init__(self, dim=8, vocab=16):
37+
super().__init__()
38+
self.model, self.embed = _TinyDecoder(dim), nn.Embedding(vocab, dim)
39+
40+
def get_input_embeddings(self):
41+
return self.embed
42+
43+
44+
class _TinyEncoder(nn.Module):
45+
def __init__(self, input_size=80, **kwargs):
46+
super().__init__()
47+
self.lin = nn.Linear(input_size, 4)
48+
49+
def output_size(self):
50+
return 4
51+
52+
53+
class _TinyAdaptor(nn.Module):
54+
def __init__(self, encoder_dim=4, llm_dim=8, **kwargs):
55+
super().__init__()
56+
self.lin = nn.Linear(encoder_dim, llm_dim)
57+
58+
59+
def _load_recipe_model():
60+
previous = tables.model_classes.get("FunASRNano")
61+
if RECIPE_DIR not in sys.path:
62+
sys.path.insert(0, RECIPE_DIR) # the recipe imports ``ctc`` and ``tools.utils`` bare
63+
path = os.path.join(RECIPE_DIR, "model.py")
64+
spec = importlib.util.spec_from_file_location("fun_asr_nano_recipe_model", path)
65+
module = importlib.util.module_from_spec(spec)
66+
sys.modules[spec.name] = module # tables.register() calls inspect.getfile on the class
67+
try:
68+
spec.loader.exec_module(module)
69+
except ImportError as exc:
70+
sys.modules.pop(spec.name, None)
71+
pytest.skip(f"recipe model.py not importable here: {exc}")
72+
finally:
73+
tables.model_classes["FunASRNano"] = previous # the package class, registered at import
74+
return module.FunASRNano
75+
76+
77+
@pytest.fixture(scope="module", params=["funasr", "recipe"])
78+
def model_class(request):
79+
return funasr_model.FunASRNano if request.param == "funasr" else _load_recipe_model()
80+
81+
82+
@pytest.fixture
83+
def build(monkeypatch, model_class):
84+
"""``build(llm_conf) -> FunASRNano`` on the CPU, with the transformers loaders stubbed."""
85+
fake = types.ModuleType("transformers")
86+
fake.AutoConfig = types.SimpleNamespace(from_pretrained=lambda path, **kw: {})
87+
fake.AutoModelForCausalLM = types.SimpleNamespace(from_config=lambda config, **kw: _TinyLLM())
88+
monkeypatch.setattr(funasr_model, "AutoConfig", fake.AutoConfig) # bound at import time
89+
monkeypatch.setattr(funasr_model, "AutoModelForCausalLM", fake.AutoModelForCausalLM)
90+
monkeypatch.setitem(sys.modules, "transformers", fake) # the recipe imports in __init__
91+
monkeypatch.setitem(tables.encoder_classes, "TinyEnc", _TinyEncoder)
92+
monkeypatch.setitem(tables.adaptor_classes, "TinyAdp", _TinyAdaptor)
93+
kw = dict(audio_encoder="TinyEnc", audio_adaptor="TinyAdp", llm="tiny")
94+
return lambda c: model_class(audio_encoder_conf={}, audio_adaptor_conf={}, llm_conf=c, **kw)
95+
96+
97+
def _run(model, batch):
98+
device = next(model.llm.parameters()).device
99+
return model.llm.model.forward(inputs_embeds=torch.zeros(batch, 3, 8, device=device))
100+
101+
102+
def test_default_leaves_sdpa_flag_and_forward_alone(build):
103+
flag = torch.backends.cuda.cudnn_sdp_enabled()
104+
model = build({})
105+
assert torch.backends.cuda.cudnn_sdp_enabled() == flag
106+
assert "forward" not in model.llm.model.__dict__ # still the class method
107+
_run(model, 2)
108+
assert model.llm.model.eager_calls[-1]["cudnn"] == flag
109+
110+
111+
def test_sdpa_backends_disable_cudnn_only_inside_forward(build):
112+
flag = torch.backends.cuda.cudnn_sdp_enabled()
113+
model = build({"sdpa_backends": ["flash", "efficient", "math"]})
114+
assert torch.backends.cuda.cudnn_sdp_enabled() == flag # construction changed nothing
115+
_run(model, 2)
116+
assert model.llm.model.eager_calls[-1]["cudnn"] is False
117+
assert torch.backends.cuda.cudnn_sdp_enabled() == flag # restored on return
118+
119+
120+
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device")
121+
def test_torch_compile_wraps_cuda_llm_and_keeps_batch_one_eager(build):
122+
model = build({"torch_compile": True}).cuda() # built on the CPU and moved, as the trainer does
123+
decoder = model.llm.model
124+
assert "forward" in decoder.__dict__
125+
_run(model, 1)
126+
assert [c["batch"] for c in decoder.eager_calls] == [1]
127+
out = _run(model, 2) # compiled: the stand-in does not record under Dynamo
128+
assert [c["batch"] for c in decoder.eager_calls] == [1]
129+
torch.testing.assert_close(out, decoder.proj(torch.zeros(2, 3, 8, device=out.device)))
130+
131+
132+
def test_torch_compile_keeps_cpu_decoder_eager(build):
133+
model = build({"torch_compile": True})
134+
_run(model, 1)
135+
_run(model, 4)
136+
assert [c["batch"] for c in model.llm.model.eager_calls] == [1, 4]

0 commit comments

Comments
 (0)