Fine-tune a transformer on mps?

[adamkgoldfarb]

Hello!

I've been using the tok2vec ensemble as the embedding layer for a multi-label non-exclusive textcat model. It's now much faster to train on m-series CPUs, so thanks for thinc-apple-ops! I have naturally become curious about the potential for improved model performance via transformers.

My early experiments incorporating distilroberta-base are very promising, and I'd like to experiment with other models but do not have access to a GPU other than the M2 GPU. I understand that currently only trf inference is supported on Apple GPUs through spaCy, but one can now train using Apple GPUs through PyTorch-- at least, most of the ops are supported and some rely on CPU fallback.

Is the best course of action fine-tuning a trf model on my data using Apple GPUs through huggingface (with a classification head, which will be discarded when the model is sourced in the spacy config), wrapping it, and freezing it to train the textcat component with cpu? Are there other setups you might recommend?

Since this isn't urgent I'm happy to wait for Apple GPU training support directly via spaCy but understand that could be a ways away, but if that's not at all on your roadmap I could investigate some other supported solution like the one above. I'd like to stay in spaCy's ecosystem as much as possible since the projects and config system are so nice!

Thanks for all you do!

[danieldk]

I understand that currently only trf inference is supported on Apple GPUs through spaCy, but one can now train using Apple GPUs through PyTorch-- at least, most of the ops are supported and some rely on CPU fallback.

spaCy itself supports training on Apple Silicon GPUs. The issue was that very recently Torch still had some bugs that lead to crashes during training. There is no fundamental difference here between spaCy and Huggingface transformers, because the spacy-transformers package uses Huggingface transformers.

These issues may be resolved now and it may depend on the model (since there is some variety in the ops used by models). For example, I just tested finetuning of our German transformer pipeline, which uses the German BERT, and it works fine. Even without CPU fallbacks.

So, I'd encourage you to just try and see if it works. If you are running into issues with distilroberta-base, we can see if we can isolate the problem and report the bug to the Torch issue tracker.

[adamkgoldfarb]

Thanks very much for looking into the issue on German BERT!

I was trying out "microsoft/deberta-v3-xsmall" and got this bug:

TypeError: Operation 'sign_out_mps()' does not support input type 'int64' in MPS backend.

This op is on the tracking list so it'll get there eventually.

I tried out distilroberta-base and it worked, but much much faster on CPU than GPU-- I had not tried on GPU before due to the bug above that I assumed would apply across Huggingface models, but glad to be wrong! I the speed issue might be addressed in torch 1.14 based on this thread, but that remains to be seen-- I may try the nightly install if I have the time.

Regardless, I appreciate your attention here and will track pytorch on mps, hoping for improved training speeds over CPU as time goes on.

Edit for posterity-- here's the config I've been using:

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

[system]
gpu_allocator = "pytorch"
seed = 0

[nlp]
lang = "en"
pipeline = ["transformer","textcat_multilabel"]
batch_size = 128
disabled = []
before_creation = null
after_creation = null
after_pipeline_creation = null
tokenizer = {"@tokenizers":"spacy.Tokenizer.v1"}

[components]

[components.textcat_multilabel]
factory = "textcat_multilabel"
scorer = {"@scorers":"spacy.textcat_multilabel_scorer.v2"}
threshold = 0.5

[components.textcat_multilabel.model]
@architectures = "spacy.TextCatEnsemble.v2"
nO = null

[components.textcat_multilabel.model.linear_model]
@architectures = "spacy.TextCatBOW.v2"
exclusive_classes = false
ngram_size = 1
no_output_layer = false
nO = null

[components.textcat_multilabel.model.tok2vec]
@architectures = "spacy-transformers.TransformerListener.v1"
grad_factor = 1.0
pooling = {"@layers":"reduce_mean.v1"}
upstream = "*"

[components.transformer]
factory = "transformer"
max_batch_items = 4096
set_extra_annotations = {"@annotation_setters":"spacy-transformers.null_annotation_setter.v1"}

[components.transformer.model]
@architectures = "spacy-transformers.TransformerModel.v3"
name = "microsoft/deberta-v3-xsmall"
# name = "distilroberta-base"
mixed_precision = false

[components.transformer.model.get_spans]
@span_getters = "spacy-transformers.strided_spans.v1"
window = 128
stride = 96

[components.transformer.model.grad_scaler_config]

[components.transformer.model.tokenizer_config]
use_fast = true

[components.transformer.model.transformer_config]

[corpora]

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

[corpora.train]
@readers = "spacy.Corpus.v1"
path = ${paths.train}
max_length = 500
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 = 3200
max_epochs = 2
max_steps = 0
eval_frequency = 400
frozen_components = []
annotating_components = []
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.v3"
progress_bar = "eval"
console_output = true
output_file = null

[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]
cats_score = 1.0
cats_score_desc = null
cats_micro_p = null
cats_micro_r = null
cats_micro_f = null
cats_macro_p = null
cats_macro_r = null
cats_macro_f = null
cats_macro_auc = null
cats_f_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.components.textcat_multilabel]

[initialize.components.textcat_multilabel.labels]
@readers = "spacy.read_labels.v1"
path = "corpus/eric/eric_education_level_labels.json"

[initialize.tokenizer]