NER model fine-tuning error with a frozen transformer

[zundamo-ti]

I'm trying to fine-tune an NER model on a custom dataset.

My pipeline includes transformer and ner components, both of which are sourced from ja_core_news_trf.

When I included the transformer component in frozen_components during training, I encountered the following error:

Saving to output directory: models/config_trf
Using GPU: True
=========================== Initializing pipeline ===========================
✔ Initialized pipeline
============================= Training pipeline =============================
ℹ Pipeline: ['transformer', 'ner']
ℹFrozen components: ['transformer']
ℹ Set annotations on update for: ['transformer']
ℹ Initial learn rate: 0.0
E    #       LOSS NER  ENTS_F  ENTS_P  ENTS_R  SCORE 
---  ------  --------  ------  ------  ------  ------
⚠ Aborting and saving the final best model. Encountered exception:
AssertionError()

Does anyone know what might be causing this issue? Any advice on my training configuration or script would be greatly appreciated.

Full Stack Trace of The Error

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[1], line 6
      4 config_path = "./configs/config_trf.cfg"
      5 output_path = "./models/config_trf"
----> 6 train(config_path, output_path, use_gpu=True, overrides={
      7     "paths.train": "./corpus/en_label/train.spacy",
      8     "paths.dev": "./corpus/en_label/dev.spacy",
      9 })

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/spacy/cli/train.py:84, in train(config_path, output_path, use_gpu, overrides)
     82 msg.good("Initialized pipeline")
     83 msg.divider("Training pipeline")
---> 84 train_nlp(nlp, output_path, use_gpu=use_gpu, stdout=sys.stdout, stderr=sys.stderr)

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/spacy/training/loop.py:135, in train(nlp, output_path, use_gpu, stdout, stderr)
    127     if output_path is not None:
    128         stdout.write(
    129             msg.warn(
    130                 f"Aborting and saving the final best model. "
   (...)
    133             + "\n"
    134         )
--> 135     raise e
    136 finally:
    137     finalize_logger()

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/spacy/training/loop.py:118, in train(nlp, output_path, use_gpu, stdout, stderr)
    116     log_step, finalize_logger = train_logger(nlp, stdout, stderr)
    117 try:
--> 118     for batch, info, is_best_checkpoint in training_step_iterator:
    119         if is_best_checkpoint is not None:
    120             with nlp.select_pipes(disable=frozen_components):

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/spacy/training/loop.py:220, in train_while_improving(nlp, optimizer, train_data, evaluate, dropout, eval_frequency, accumulate_gradient, patience, max_steps, exclude, annotating_components, before_update)
    218 dropout = next(dropouts)  # type: ignore
    219 for subbatch in subdivide_batch(batch, accumulate_gradient):
--> 220     nlp.update(
    221         subbatch,
    222         drop=dropout,
    223         losses=losses,
    224         sgd=False,  # type: ignore[arg-type]
    225         exclude=exclude,
    226         annotates=annotating_components,
    227     )
    228 # TODO: refactor this so we don't have to run it separately in here
    229 for name, proc in nlp.pipeline:

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/spacy/language.py:1197, in Language.update(self, examples, _, drop, sgd, losses, component_cfg, exclude, annotates)
   1194 for name, proc in self.pipeline:
   1195     # ignore statements are used here because mypy ignores hasattr
   1196     if name not in exclude and hasattr(proc, "update"):
-> 1197         proc.update(examples, sgd=None, losses=losses, **component_cfg[name])  # type: ignore
   1198     if sgd not in (None, False):
   1199         if (
   1200             name not in exclude
   1201             and isinstance(proc, ty.TrainableComponent)
   1202             and proc.is_trainable
   1203             and proc.model not in (True, False, None)
   1204         ):

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/spacy/pipeline/transition_parser.pyx:419, in spacy.pipeline.transition_parser.Parser.update()

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/thinc/model.py:328, in Model.begin_update(self, X)
    321 def begin_update(self, X: InT) -> Tuple[OutT, Callable[[OutT], InT]]:
    322     """Run the model over a batch of data, returning the output and a
    323     callback to complete the backward pass. A tuple (Y, finish_update),
    324     where Y is a batch of output data, and finish_update is a callback that
    325     takes the gradient with respect to the output and an optimizer function,
    326     and returns the gradient with respect to the input.
    327     """
--> 328     return self._func(self, X, is_train=True)

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/spacy/ml/tb_framework.py:33, in forward(model, X, is_train)
     32 def forward(model, X, is_train):
---> 33     step_model = ParserStepModel(
     34         X,
     35         model.layers,
     36         unseen_classes=model.attrs["unseen_classes"],
     37         train=is_train,
     38         has_upper=model.attrs["has_upper"],
     39     )
     41     return step_model, step_model.finish_steps

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/spacy/ml/parser_model.pyx:250, in spacy.ml.parser_model.ParserStepModel.__init__()

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/thinc/model.py:310, in Model.__call__(self, X, is_train)
    307 def __call__(self, X: InT, is_train: bool) -> Tuple[OutT, Callable]:
    308     """Call the model's `forward` function, returning the output and a
    309     callback to compute the gradients via backpropagation."""
--> 310     return self._func(self, X, is_train=is_train)

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/thinc/layers/chain.py:54, in forward(model, X, is_train)
     52 callbacks = []
     53 for layer in model.layers:
---> 54     Y, inc_layer_grad = layer(X, is_train=is_train)
     55     callbacks.append(inc_layer_grad)
     56     X = Y

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/thinc/model.py:310, in Model.__call__(self, X, is_train)
    307 def __call__(self, X: InT, is_train: bool) -> Tuple[OutT, Callable]:
    308     """Call the model's `forward` function, returning the output and a
    309     callback to compute the gradients via backpropagation."""
--> 310     return self._func(self, X, is_train=is_train)

File ~/Playground/ner-ja-wikipedia/.venv/lib/python3.10/site-packages/spacy_curated_transformers/models/listeners.py:392, in last_transformer_layer_listener_forward(model, docs, is_train)
    389 _backprop = ListenerStateUtils.get_backprop(model)
    391 if is_train:
--> 392     assert _outputs is not None
    393     ListenerStateUtils.verify_inputs(model, docs)
    395     Y, backprop_pooling = pooling(_outputs.last_hidden_layer_states, is_train)

AssertionError:

My Training Script

from spacy.cli.train import train


config_path = "./configs/config_trf.cfg"
output_path = "./models/config_trf"
train(config_path, output_path, use_gpu=True, overrides={
    "paths.train": "./corpus/en_label/train.spacy",
    "paths.dev": "./corpus/en_label/dev.spacy",
})

My Config

[paths]
train = null
dev = null
vectors = null
init_tok2vec = null

[system]
gpu_allocator = "pytorch"
seed = 0

[nlp]
lang = "ja"
pipeline = ["transformer","ner"]
batch_size = 128
disabled = []
before_creation = null
after_creation = null
after_pipeline_creation = null
vectors = {"@vectors":"spacy.Vectors.v1"}

[nlp.tokenizer]
@tokenizers = "spacy.ja.JapaneseTokenizer"
split_mode = null

[components]

[components.ner]
source = "ja_core_news_trf"

[components.transformer]
source = "ja_core_news_trf"

[corpora]

[corpora.dev]
@readers = "spacy.Corpus.v1"
path = ${paths.dev}
max_length = 0
gold_preproc = false
limit = 0
augmenter = null

[corpora.train]
@readers = "spacy.Corpus.v1"
path = ${paths.train}
max_length = 0
gold_preproc = false
limit = 0
augmenter = null

[training]
accumulate_gradient = 3
dev_corpus = "corpora.dev"
train_corpus = "corpora.train"
seed = ${system.seed}
gpu_allocator = ${system.gpu_allocator}
dropout = 0.1
patience = 1600
max_epochs = 0
max_steps = 20000
eval_frequency = 200
frozen_components = ["transformer"]
annotating_components = ["transformer"]
before_to_disk = null
before_update = null

[training.batcher]
@batchers = "spacy.batch_by_padded.v1"
discard_oversize = true
size = 2000
buffer = 256
get_length = null

[training.logger]
@loggers = "spacy.ConsoleLogger.v1"
progress_bar = false

[training.optimizer]
@optimizers = "Adam.v1"
beta1 = 0.9
beta2 = 0.999
L2_is_weight_decay = true
L2 = 0.01
grad_clip = 1.0
use_averages = false
eps = 0.00000001

[training.optimizer.learn_rate]
@schedules = "warmup_linear.v1"
warmup_steps = 250
total_steps = 20000
initial_rate = 0.00005

[training.score_weights]
ents_f = 1.0
ents_p = 0.0
ents_r = 0.0
ents_per_type = null

[pretraining]

[initialize]
vectors = ${paths.vectors}
init_tok2vec = ${paths.init_tok2vec}
vocab_data = null
lookups = null
before_init = null
after_init = null

[initialize.components]

[initialize.tokenizer]

My Environment

$ spacy info

============================== Info about spaCy ==============================

spaCy version    3.8.7                         
Location         /path/to/my-project/.venv/lib/python3.10/site-packages/spacy
Platform         macOS-15.3.2-arm64-arm-64bit  
Python version   3.10.15                       
Pipelines        ja_core_news_lg (3.8.0), ja_core_news_trf (3.8.0)

Similar question: #9233 (comment)